chore: import upstream snapshot with attribution

python/ray/autoscaler/v2/tests/conftest.py

@@ -0,0 +1 @@
+from ray.tests.conftest import *  # noqa

python/ray/autoscaler/v2/tests/test_allocation_timeout_recovery.py

@@ -0,0 +1,544 @@
+# coding: utf-8
+"""
+Test recovery behavior for the ALLOCATION_TIMEOUT scenario.
+
+Verifies that when an instance reaches ALLOCATION_TIMEOUT, the autoscaler can:
+1. Terminate the timed-out old instance.
+2. Launch a replacement instance.
+3. Avoid a QUEUED->REQUESTED->QUEUED loop.
+
+Test design principles:
+- Pure Python, mocking only the k8s client.
+- Validate instance state rather than log output.
+- Run reconcile multiple times to verify the state does not get stuck.
+"""
+
+import time
+from typing import Any, Dict, List
+
+import pytest
+
+from ray.autoscaler.v2.instance_manager.config import InstanceReconcileConfig, Provider
+from ray.autoscaler.v2.instance_manager.instance_manager import InstanceManager
+from ray.autoscaler.v2.instance_manager.instance_storage import InstanceStorage
+from ray.autoscaler.v2.instance_manager.node_provider import (
+    CloudInstance,
+    ICloudInstanceProvider,
+)
+from ray.autoscaler.v2.instance_manager.reconciler import Reconciler
+from ray.autoscaler.v2.instance_manager.storage import InMemoryStorage
+from ray.autoscaler.v2.instance_manager.subscribers.cloud_instance_updater import (
+    CloudInstanceUpdater,
+)
+from ray.autoscaler.v2.instance_manager.subscribers.cloud_resource_monitor import (
+    CloudResourceMonitor,
+)
+from ray.autoscaler.v2.scheduler import (
+    ResourceDemandScheduler,
+)
+from ray.autoscaler.v2.tests.util import create_instance
+from ray.core.generated.autoscaler_pb2 import (
+    ClusterResourceState,
+    NodeState,
+    NodeStatus,
+)
+from ray.core.generated.instance_manager_pb2 import Instance, NodeKind
+
+s_to_ns = 1 * 1_000_000_000
+
+
+class MockAutoscalingConfig:
+    """Mock autoscaling config for testing"""
+
+    def __init__(self, configs=None):
+        if configs is None:
+            configs = {}
+        self._configs = configs
+
+    def get_node_type_configs(self):
+        return self._configs.get("node_type_configs", {})
+
+    def get_max_num_worker_nodes(self):
+        return self._configs.get("max_num_worker_nodes")
+
+    def get_max_num_nodes(self):
+        n = self._configs.get("max_num_worker_nodes")
+        return n + 1 if n is not None else None
+
+    def get_upscaling_speed(self):
+        return self._configs.get("upscaling_speed", 0.0)
+
+    def get_max_concurrent_launches(self):
+        return self._configs.get("max_concurrent_launches", 100)
+
+    def get_instance_reconcile_config(self):
+        return self._configs.get("instance_reconcile_config", InstanceReconcileConfig())
+
+    def disable_node_updaters(self):
+        return self._configs.get("disable_node_updaters", True)
+
+    def disable_launch_config_check(self):
+        return self._configs.get("disable_launch_config_check", False)
+
+    def get_idle_timeout_s(self):
+        return self._configs.get("idle_timeout_s", 999)
+
+    def get_provider_instance_type(self, ray_node_type):
+        return ""
+
+    @property
+    def provider(self):
+        return Provider.UNKNOWN
+
+
+class EventCapturingSubscriber:
+    """
+    Subscriber that captures events for verification and delegates to CloudInstanceUpdater.
+
+    This wraps the real CloudInstanceUpdater to capture events for test assertions
+    while using the actual launch/terminate logic.
+    """
+
+    def __init__(self, cloud_provider: ICloudInstanceProvider):
+        self.cloud_provider = cloud_provider
+        self.updater = CloudInstanceUpdater(cloud_provider=cloud_provider)
+        self.events = []
+
+    def notify(self, events):
+        self.events.extend(events)
+        # Delegate to real updater which calls cloud_provider.launch/terminate
+        self.updater.notify(events)
+
+    def clear(self):
+        self.events.clear()
+
+    def events_by_id(self, instance_id):
+        return [e for e in self.events if e.instance_id == instance_id]
+
+
+class MockK8sClient:
+    """
+    Mock Kubernetes API client that simulates RayCluster behavior.
+
+    Tracks:
+    - replicas: current replica count
+    - workers_to_delete: list of worker pod names to delete
+    - patch_history: history of all patches for verification
+    """
+
+    def __init__(self, max_replicas=3, initial_replicas=3, worker_pod_names=None):
+        self.max_replicas = max_replicas
+        self.replicas = initial_replicas
+        self.workers_to_delete: List[str] = []
+        self.patch_history: List[Dict] = []
+        self._resource_version = 100
+        # Worker pod names - defaults to worker-001, worker-002, worker-003
+        self.worker_pod_names = worker_pod_names or [
+            "worker-001",
+            "worker-002",
+            "worker-003",
+        ]
+
+    def get(self, path: str) -> Dict[str, Any]:
+        """Handle GET requests"""
+        if "rayclusters" in path:
+            return {
+                "metadata": {
+                    "name": "test-ray-cluster",
+                    "namespace": "default",
+                    "resourceVersion": str(self._resource_version),
+                },
+                "spec": {
+                    "workerGroupSpecs": [
+                        {
+                            "groupName": "default-worker-group",
+                            "replicas": self.replicas,
+                            "minReplicas": 0,
+                            "maxReplicas": self.max_replicas,
+                            "scaleStrategy": {
+                                "workersToDelete": list(self.workers_to_delete)
+                            },
+                            "numOfHosts": 1,
+                        }
+                    ]
+                },
+            }
+        elif "pods" in path:
+            # Return pods based on current replicas (excluding workers_to_delete)
+            # Use worker_pod_names to match the cloud_instances in the test
+            items = []
+            for i in range(min(self.replicas, len(self.worker_pod_names))):
+                pod_name = self.worker_pod_names[i]
+                if pod_name not in self.workers_to_delete:
+                    # worker-003 is in ALLOCATION_TIMEOUT state, so it should not be running
+                    if pod_name == "worker-003":
+                        # Pod is pending/failed - not running
+                        container_state = {"waiting": {"reason": "ContainerCreating"}}
+                    else:
+                        container_state = {"running": {}}
+
+                    items.append(
+                        {
+                            "metadata": {
+                                "name": pod_name,
+                                "labels": {
+                                    "ray.io/cluster": "test-ray-cluster",
+                                    "ray.io/node-type": "worker",
+                                    "ray.io/group": "default-worker-group",
+                                },
+                            },
+                            "status": {
+                                "containerStatuses": [{"state": container_state}]
+                            },
+                        }
+                    )
+            return {
+                "metadata": {"resourceVersion": str(self._resource_version)},
+                "items": items,
+            }
+        return {}
+
+    def patch(self, path: str, payload: List[Dict]) -> Dict[str, Any]:
+        """Handle PATCH requests and update internal state"""
+        self.patch_history.append({"path": path, "payload": payload})
+        self._resource_version += 1
+
+        for op in payload:
+            if op["op"] == "replace" and "replicas" in op["path"]:
+                self.replicas = op["value"]
+            elif op["op"] == "replace" and "scaleStrategy" in op["path"]:
+                self.workers_to_delete = op["value"].get("workersToDelete", [])
+
+        return {}
+
+
+class TestAllocationTimeoutRecovery:
+    """Test ALLOCATION_TIMEOUT instance recovery"""
+
+    @staticmethod
+    def _add_instances(instance_storage, instances):
+        for instance in instances:
+            ok, _ = instance_storage.upsert_instance(instance)
+            assert ok
+
+    def test_no_queued_requested_loop(self):
+        """
+        Minimal reproduction path:
+
+        Preconditions:
+        - maxReplicas = 3
+        - 3 worker instances: 2 ALLOCATED (healthy), 1 ALLOCATION_TIMEOUT
+        - idle_worker_nodes = 1
+
+        Steps:
+        1. Run Reconciler.reconcile() multiple times to simulate repeated
+           reconciler cycles.
+
+        Expected results:
+        1. The new instance does not enter a QUEUED->REQUESTED->QUEUED loop.
+        2. The new instance can eventually transition into REQUESTED.
+        3. Terminate is submitted to Kubernetes before launch.
+        """
+        # ===== Preconditions =====
+        instance_storage = InstanceStorage(
+            cluster_id="test_cluster_id",
+            storage=InMemoryStorage(),
+        )
+
+        cloud_resource_monitor = CloudResourceMonitor()
+
+        # Mock K8s client: maxReplicas=3, initial_replicas=3
+        # Worker pod names match the cloud_instance_id in cloud_instances
+        mock_k8s = MockK8sClient(
+            max_replicas=3,
+            initial_replicas=3,
+            worker_pod_names=["worker-001", "worker-002", "worker-003"],
+        )
+
+        # Create real cloud provider with mock k8s client
+        from unittest.mock import MagicMock
+
+        from ray.autoscaler.v2.instance_manager.cloud_providers.kuberay.cloud_provider import (
+            KubeRayProvider,
+        )
+
+        cloud_provider = KubeRayProvider(
+            cluster_name="test-ray-cluster",
+            provider_config={"namespace": "default"},
+            gcs_client=MagicMock(),
+            k8s_api_client=mock_k8s,
+        )
+
+        # Use EventCapturingSubscriber that wraps real CloudInstanceUpdater
+        # and captures events for test verification
+        mock_subscriber = EventCapturingSubscriber(cloud_provider=cloud_provider)
+
+        instance_manager = InstanceManager(
+            instance_storage=instance_storage,
+            instance_status_update_subscribers=[mock_subscriber],
+        )
+
+        # Create 2 ALLOCATED instances and 1 ALLOCATION_TIMEOUT instance.
+        current_time = time.time_ns()
+        timeout_time = (
+            current_time - 200 * s_to_ns
+        )  # 200s ago, beyond the timeout threshold.
+
+        from ray._common.utils import binary_to_hex
+
+        instances = [
+            # Head node
+            create_instance(
+                "head",
+                status=Instance.RAY_RUNNING,
+                cloud_instance_id="head-001",
+                node_kind=NodeKind.HEAD,
+                instance_type="head",
+                ray_node_id=binary_to_hex(b"head"),
+            ),
+            # Worker 1: healthy
+            create_instance(
+                "worker-1",
+                status=Instance.ALLOCATED,
+                instance_type="default-worker-group",
+                cloud_instance_id="worker-001",
+                node_kind=NodeKind.WORKER,
+                ray_node_id=binary_to_hex(b"wkr1"),
+            ),
+            # Worker 2: healthy
+            create_instance(
+                "worker-2",
+                status=Instance.ALLOCATED,
+                instance_type="default-worker-group",
+                cloud_instance_id="worker-002",
+                node_kind=NodeKind.WORKER,
+                ray_node_id=binary_to_hex(b"wkr2"),
+            ),
+            # Worker 3: ALLOCATION_TIMEOUT (startup timed out)
+            create_instance(
+                "worker-3",
+                status=Instance.ALLOCATION_TIMEOUT,
+                instance_type="default-worker-group",
+                cloud_instance_id="worker-003",
+                node_kind=NodeKind.WORKER,
+                status_times=[(Instance.ALLOCATION_TIMEOUT, timeout_time)],
+            ),
+        ]
+
+        TestAllocationTimeoutRecovery._add_instances(instance_storage, instances)
+
+        # Mock ray nodes: head + 2 healthy workers.
+        # ray_node_id must match the instance node_id.
+        # available_resources=0 means resources are fully consumed, so scale-up is needed.
+        ray_nodes = [
+            NodeState(
+                node_id=b"head",
+                status=NodeStatus.RUNNING,
+                instance_id="head-001",
+                total_resources={"CPU": 0},
+                available_resources={"CPU": 0},
+            ),
+            NodeState(
+                node_id=b"wkr1",
+                status=NodeStatus.RUNNING,
+                instance_id="worker-001",
+                total_resources={"CPU": 4},
+                available_resources={"CPU": 0},  # Resources are fully consumed.
+            ),
+            NodeState(
+                node_id=b"wkr2",
+                status=NodeStatus.RUNNING,
+                instance_id="worker-002",
+                total_resources={"CPU": 4},
+                available_resources={"CPU": 0},  # Resources are fully consumed.
+            ),
+        ]
+
+        # Mock cloud instances
+        # worker-003 must be included because the cloud instance still exists
+        # while the instance is in ALLOCATION_TIMEOUT.
+        # is_running=False means the pod is not running normally.
+        cloud_instances = {
+            "head-001": CloudInstance("head-001", "head", True, NodeKind.HEAD),
+            "worker-001": CloudInstance(
+                "worker-001", "default-worker-group", True, NodeKind.WORKER
+            ),
+            "worker-002": CloudInstance(
+                "worker-002", "default-worker-group", True, NodeKind.WORKER
+            ),
+            "worker-003": CloudInstance(
+                "worker-003",
+                "default-worker-group",
+                False,
+                NodeKind.WORKER,  # is_running=False
+            ),
+        }
+
+        # Use the real ResourceDemandScheduler.
+
+        scheduler = ResourceDemandScheduler()
+
+        # Add pending_resource_requests to simulate resource demand.
+        # Request 4 CPU while current available CPU is 0, so one more worker is needed.
+        from ray.core.generated.autoscaler_pb2 import (
+            ResourceRequest,
+            ResourceRequestByCount,
+        )
+
+        ray_cluster_resource_state = ClusterResourceState(
+            node_states=ray_nodes,
+            pending_resource_requests=[
+                ResourceRequestByCount(
+                    request=ResourceRequest(resources_bundle={"CPU": 4}), count=1
+                ),
+            ],
+        )
+
+        # Mock autoscaling config. Use MockAutoscalingConfig to avoid schema validation.
+        # Build real NodeTypeConfig objects.
+        from ray.autoscaler.v2.instance_manager.config import NodeTypeConfig
+
+        node_type_configs = {
+            "default-worker-group": NodeTypeConfig(
+                name="default-worker-group",
+                min_worker_nodes=0,
+                max_worker_nodes=3,
+                resources={"CPU": 4},
+                labels={},
+                launch_config_hash="hash1",
+                idle_timeout_s=None,
+            ),
+            "head": NodeTypeConfig(
+                name="head",
+                min_worker_nodes=0,
+                max_worker_nodes=1,
+                resources={"CPU": 0},
+                labels={},
+                launch_config_hash="hash1",
+                idle_timeout_s=None,
+            ),
+        }
+
+        autoscaling_config = MockAutoscalingConfig(
+            configs={
+                "node_type_configs": node_type_configs,
+                "max_num_worker_nodes": 3,
+                "upscaling_speed": 1.0,
+                "max_concurrent_launches": 100,
+                "instance_reconcile_config": InstanceReconcileConfig(
+                    request_status_timeout_s=10,
+                    allocate_status_timeout_s=300,
+                ),
+                "disable_node_updaters": True,
+                "disable_launch_config_check": True,
+                "idle_timeout_s": 999,
+            }
+        )
+
+        # ===== Steps: run reconcile multiple times =====
+        # Run 3 reconcile cycles to simulate repeated reconciler execution.
+        for cycle in range(3):
+            # Clear subscriber events so each iteration is checked independently.
+            mock_subscriber.events.clear()
+
+            Reconciler.reconcile(
+                instance_manager=instance_manager,
+                scheduler=scheduler,
+                cloud_provider=cloud_provider,
+                cloud_resource_monitor=cloud_resource_monitor,
+                ray_cluster_resource_state=ray_cluster_resource_state,
+                non_terminated_cloud_instances=cloud_instances,
+                cloud_provider_errors=[],
+                ray_install_errors=[],
+                autoscaling_config=autoscaling_config,
+            )
+
+            # Fetch current instance state to ensure repeated reconcile calls do not fail.
+            instance_storage.get_instances()
+
+        # ===== Expected result checks =====
+
+        # Get final instance states.
+        all_instances, _ = instance_storage.get_instances()
+
+        # 1. Verify the ALLOCATION_TIMEOUT instance transitions to TERMINATING.
+        worker_3 = all_instances.get("worker-3")
+        assert worker_3 is not None, "Expected worker-3 instance to exist"
+        assert worker_3.status == Instance.TERMINATING, (
+            f"Expected worker-3 status to be TERMINATING, "
+            f"got {Instance.InstanceStatus.Name(worker_3.status)}"
+        )
+
+        # 2. Verify at least one new instance is created (QUEUED or REQUESTED).
+        new_instances = [
+            i
+            for i in all_instances.values()
+            if i.status in (Instance.QUEUED, Instance.REQUESTED)
+        ]
+        assert len(new_instances) >= 1, (
+            f"Expected at least 1 new instance (QUEUED/REQUESTED), "
+            f"got {len(new_instances)}"
+        )
+
+        # 3. Verify K8s patch history shows terminate before launch.
+        # Each patch may contain both replicas and scaleStrategy updates.
+        # Key checks:
+        #   a) There should be two patches (terminate + launch).
+        #   b) The first patch should contain workersToDelete.
+        #   c) The second patch should increase replicas.
+
+        # Verify there are two patches (terminate + launch).
+        assert len(mock_k8s.patch_history) == 2, (
+            f"Expected 2 patches (terminate + launch), got {len(mock_k8s.patch_history)}. "
+            "If only 1 patch, launch may have failed due to maxReplicas bug."
+        )
+
+        # Verify the first patch contains workersToDelete.
+        first_patch = mock_k8s.patch_history[0]
+        first_patch_payload_str = str(first_patch.get("payload", []))
+        assert "workersToDelete" in first_patch_payload_str, (
+            f"Expected first patch to contain workersToDelete (terminate before launch). "
+            f"First patch: {first_patch}"
+        )
+
+        # Verify workersToDelete contains the timed-out instance.
+        for op in first_patch["payload"]:
+            if "scaleStrategy" in op.get("path", ""):
+                workers_to_delete = op["value"].get("workersToDelete", [])
+                assert (
+                    "worker-003" in workers_to_delete
+                ), f"Expected worker-003 in workersToDelete, got {workers_to_delete}"
+                break
+
+        # Verify the second patch increases replicas.
+        second_patch = mock_k8s.patch_history[1]
+        for op in second_patch["payload"]:
+            if "replicas" in op.get("path", ""):
+                assert (
+                    op["value"] == 3
+                ), f"Expected replicas=3 after launch, got {op['value']}"
+                break
+
+        # 4. Verify final replicas do not exceed maxReplicas.
+        assert (
+            mock_k8s.replicas <= 3
+        ), f"Expected replicas <= 3, got {mock_k8s.replicas}"
+
+        # 5. Verify no instance enters a REQUESTED->QUEUED->REQUESTED loop.
+        # Check whether status_history contains a repeated transition pattern.
+        for instance in all_instances.values():
+            status_sequence = [h.instance_status for h in instance.status_history]
+            # Check for a REQUESTED->QUEUED->REQUESTED pattern.
+            for i in range(len(status_sequence) - 2):
+                assert not (
+                    status_sequence[i] == Instance.REQUESTED
+                    and status_sequence[i + 1] == Instance.QUEUED
+                    and status_sequence[i + 2] == Instance.REQUESTED
+                ), (
+                    f"Instance {instance.instance_id} entered REQUESTED->QUEUED->REQUESTED loop. "
+                    f"Status sequence: {[Instance.InstanceStatus.Name(s) for s in status_sequence]}"
+                )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "-s"])

python/ray/autoscaler/v2/tests/test_autoscaler.py

@@ -0,0 +1,262 @@
+import logging
+import os
+import subprocess
+import sys
+import time
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+import ray
+from ray._common.test_utils import wait_for_condition
+from ray._raylet import GcsClient
+from ray.autoscaler._private.fake_multi_node.node_provider import FAKE_HEAD_NODE_ID
+from ray.autoscaler.v2.autoscaler import Autoscaler
+from ray.autoscaler.v2.event_logger import AutoscalerEventLogger
+from ray.autoscaler.v2.instance_manager.config import AutoscalingConfig
+from ray.autoscaler.v2.monitor import AutoscalerMonitor
+from ray.autoscaler.v2.sdk import get_cluster_status, request_cluster_resources
+from ray.autoscaler.v2.tests.util import MockEventLogger
+from ray.cluster_utils import Cluster
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_AUTOSCALING_CONFIG = {
+    "cluster_name": "fake_multinode",
+    "max_workers": 8,
+    "provider": {
+        "type": "fake_multinode",
+    },
+    "available_node_types": {
+        "ray.head.default": {
+            "resources": {
+                "CPU": 0,
+            },
+            "max_workers": 0,
+            "node_config": {},
+        },
+        "ray.worker.cpu": {
+            "resources": {
+                "CPU": 1,
+            },
+            "min_workers": 0,
+            "max_workers": 10,
+            "node_config": {},
+        },
+        "ray.worker.gpu": {
+            "resources": {
+                "GPU": 1,
+            },
+            "min_workers": 0,
+            "max_workers": 10,
+            "node_config": {},
+        },
+    },
+    "head_node_type": "ray.head.default",
+    "upscaling_speed": 0,
+    "idle_timeout_minutes": 0.08,  # ~5 seconds
+}
+
+
+@pytest.fixture(scope="function")
+def make_autoscaler():
+    ctx = {}
+
+    def _make_autoscaler(config):
+        head_node_kwargs = {
+            "env_vars": {
+                "RAY_CLOUD_INSTANCE_ID": FAKE_HEAD_NODE_ID,
+                "RAY_OVERRIDE_NODE_ID_FOR_TESTING": FAKE_HEAD_NODE_ID,
+                "RAY_NODE_TYPE_NAME": "ray.head.default",
+            },
+            "num_cpus": config["available_node_types"]["ray.head.default"]["resources"][
+                "CPU"
+            ],
+        }
+        cluster = Cluster(
+            initialize_head=True, head_node_args=head_node_kwargs, connect=True
+        )
+        ctx["cluster"] = cluster
+
+        mock_config_reader = MagicMock()
+        gcs_address = cluster.address
+
+        # Configs for the node provider
+        config["provider"]["gcs_address"] = gcs_address
+        config["provider"]["head_node_id"] = FAKE_HEAD_NODE_ID
+        config["provider"]["launch_multiple"] = True
+        os.environ["RAY_FAKE_CLUSTER"] = "1"
+        mock_config_reader.get_cached_autoscaling_config.return_value = (
+            AutoscalingConfig(configs=config, skip_content_hash=True)
+        )
+        gcs_address = gcs_address
+        gcs_client = GcsClient(gcs_address)
+
+        event_logger = AutoscalerEventLogger(MockEventLogger(logger))
+
+        autoscaler = Autoscaler(
+            session_name="test",
+            config_reader=mock_config_reader,
+            gcs_client=gcs_client,
+            event_logger=event_logger,
+        )
+
+        return autoscaler
+
+    yield _make_autoscaler
+    try:
+        ray.shutdown()
+        ctx["cluster"].shutdown()
+    except Exception:
+        logger.exception("Error during teardown")
+        # Run ray stop to clean up everything
+        subprocess.run(
+            ["ray", "stop", "--force"], stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        )
+
+
+def test_basic_scaling(make_autoscaler):
+    config = DEFAULT_AUTOSCALING_CONFIG
+    autoscaler = make_autoscaler(DEFAULT_AUTOSCALING_CONFIG)
+    gcs_address = autoscaler._gcs_client.address
+
+    # Resource requests
+    print("=================== Test scaling up constraint 1/2====================")
+    request_cluster_resources(
+        gcs_address, [{"resources": {"CPU": 1}}, {"resources": {"GPU": 1}}]
+    )
+
+    def verify():
+        autoscaler.update_autoscaling_state()
+        cluster_state = get_cluster_status(gcs_address)
+        assert len(cluster_state.active_nodes + cluster_state.idle_nodes) == 3
+        return True
+
+    wait_for_condition(verify, retry_interval_ms=2000)
+
+    # Test scaling down shouldn't happen
+    print("=================== Test scaling down constraint 2/2 ====================")
+
+    idle_timeout_s = config["idle_timeout_minutes"] * 60
+    time.sleep(idle_timeout_s + 2)
+
+    wait_for_condition(verify, retry_interval_ms=2000)
+
+    # Test scaling down.
+    print("=================== Test scaling down idle ====================")
+    request_cluster_resources(gcs_address, [])
+
+    def verify():
+        autoscaler.update_autoscaling_state()
+        cluster_state = get_cluster_status(gcs_address)
+        assert len(cluster_state.active_nodes + cluster_state.idle_nodes) == 1
+        return True
+
+    wait_for_condition(verify, retry_interval_ms=2000)
+
+    # Test scaling up again with tasks
+    print("=================== Test scaling up with tasks ====================")
+
+    @ray.remote
+    def task():
+        time.sleep(999)
+
+    task.options(num_cpus=1).remote()
+    task.options(num_cpus=0, num_gpus=1).remote()
+
+    def verify():
+        autoscaler.update_autoscaling_state()
+        cluster_state = get_cluster_status(gcs_address)
+        assert len(cluster_state.active_nodes + cluster_state.idle_nodes) == 3
+        return True
+
+    wait_for_condition(verify, retry_interval_ms=2000)
+
+    # Test with placement groups
+    print(
+        "=================== Test scaling up with placement groups ===================="
+    )
+
+    # Spread to create another 2 nodes.
+    ray.util.placement_group(
+        name="pg", strategy="STRICT_SPREAD", bundles=[{"CPU": 0.5}, {"CPU": 0.5}]
+    )
+
+    def verify():
+        autoscaler.update_autoscaling_state()
+        cluster_state = get_cluster_status(gcs_address)
+        assert len(cluster_state.active_nodes + cluster_state.idle_nodes) == 5
+        return True
+
+    wait_for_condition(verify, retry_interval_ms=2000)
+    # Pack with feasible ones
+    ray.util.placement_group(
+        name="pg_feasible", strategy="STRICT_PACK", bundles=[{"CPU": 0.5}, {"CPU": 0.5}]
+    )
+
+    def verify():
+        autoscaler.update_autoscaling_state()
+        cluster_state = get_cluster_status(gcs_address)
+        assert len(cluster_state.active_nodes + cluster_state.idle_nodes) == 6
+        return True
+
+    wait_for_condition(verify, retry_interval_ms=2000)
+
+    # Pack with infeasible request
+    ray.util.placement_group(
+        name="pg_infeasible", strategy="STRICT_PACK", bundles=[{"CPU": 1}, {"CPU": 1}]
+    )
+
+    def verify():
+        autoscaling_state = autoscaler.update_autoscaling_state()
+        cluster_state = get_cluster_status(gcs_address)
+        assert len(cluster_state.active_nodes + cluster_state.idle_nodes) == 6
+        assert len(autoscaling_state.infeasible_gang_resource_requests) == 1
+
+        return True
+
+    wait_for_condition(verify, retry_interval_ms=2000)
+
+    # Test report autoscaling state
+    def verify():
+        autoscaling_state = autoscaler.update_autoscaling_state()
+        assert len(autoscaling_state.infeasible_gang_resource_requests) == 1
+        # For now, we just track that it's called ok.
+        autoscaler._gcs_client.report_autoscaling_state(
+            autoscaling_state.SerializeToString()
+        )
+        return True
+
+    wait_for_condition(verify, retry_interval_ms=2000)
+
+
+class TestAutoscalerMonitor(AutoscalerMonitor):
+    """Lightweight wrapper for testing _run() without full init."""
+
+    def __init__(self, gcs_address, gcs_client, autoscaler):
+        self.gcs_address = gcs_address
+        self.gcs_client = gcs_client
+        self.autoscaler = autoscaler
+        self._session_name = "test"
+
+
+def test_raise_AuthenticationError_v2(make_autoscaler):
+    autoscaler = make_autoscaler(DEFAULT_AUTOSCALING_CONFIG)
+    gcs_client = autoscaler._gcs_client
+    gcs_address = gcs_client.address
+
+    monitor = TestAutoscalerMonitor(gcs_address, gcs_client, autoscaler)
+
+    def flaky():
+        raise ray.exceptions.AuthenticationError("WrongClusterID")
+
+    with patch.object(autoscaler, "update_autoscaling_state", side_effect=flaky):
+        with pytest.raises(ray.exceptions.AuthenticationError):
+            monitor._run()
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/test_config.py

@@ -0,0 +1,251 @@
+# coding: utf-8
+import os
+import sys
+import tempfile
+
+import pytest  # noqa
+
+from ray._common.utils import binary_to_hex
+from ray._private.test_utils import get_test_config_path
+from ray.autoscaler import AUTOSCALER_DIR_PATH
+from ray.autoscaler._private.util import format_readonly_node_type
+from ray.autoscaler.v2.instance_manager import config as config_mod
+from ray.autoscaler.v2.instance_manager.config import (
+    FileConfigReader,
+    Provider,
+    ReadOnlyProviderConfigReader,
+)
+
+
+@pytest.mark.parametrize(
+    "skip_hash",
+    [True, False],
+)
+def test_simple(skip_hash):
+    config = FileConfigReader(
+        get_test_config_path("test_multi_node.yaml"), skip_content_hash=skip_hash
+    ).get_cached_autoscaling_config()
+    assert config.get_cloud_node_config("head_node") == {"InstanceType": "m5.large"}
+    assert config.get_docker_config("head_node") == {
+        "image": "anyscale/ray-ml:latest",
+        "container_name": "ray_container",
+        "pull_before_run": True,
+    }
+    assert config.get_worker_start_ray_commands
+
+
+def test_complex():
+    config = FileConfigReader(
+        get_test_config_path("test_ray_complex.yaml")
+    ).get_cached_autoscaling_config()
+    assert config.get_head_setup_commands() == [
+        "echo a",
+        "echo b",
+        "echo ${echo hi}",
+        "echo head",
+    ]
+    assert config.get_head_start_ray_commands() == [
+        "ray stop",
+        "ray start --head --autoscaling-config=~/ray_bootstrap_config.yaml",
+    ]
+    assert config.get_worker_setup_commands("worker_nodes") == [
+        "echo a",
+        "echo b",
+        "echo ${echo hi}",
+        "echo worker",
+    ]
+    assert config.get_worker_start_ray_commands() == [
+        "ray stop",
+        "ray start --address=$RAY_HEAD_IP",
+    ]
+    assert config.get_worker_setup_commands("worker_nodes1") == [
+        "echo worker1",
+    ]
+
+    assert config.get_docker_config("head_node") == {
+        "image": "anyscale/ray-ml:head-default",
+        "container_name": "ray_container",
+        "pull_before_run": True,
+    }
+
+    assert config.get_docker_config("default") == {
+        "image": "anyscale/ray-ml:worker-default",
+        "container_name": "ray_container",
+        "pull_before_run": True,
+    }
+
+    assert config.get_docker_config("worker_nodes") == {
+        "image": "anyscale/ray-ml:worker-default",
+        "container_name": "ray_container",
+        "pull_before_run": True,
+    }
+
+    assert config.get_docker_config("worker_nodes1") == {
+        "image": "anyscale/ray-ml:worker_nodes1",
+        "container_name": "ray_container",
+        "pull_before_run": True,
+    }
+
+    assert config.get_initialization_commands("worker_nodes") == ["echo what"]
+
+    assert config.get_initialization_commands("worker_nodes1") == ["echo init"]
+
+    assert config.get_node_resources("worker_nodes1") == {"CPU": 2}
+
+    assert config.get_node_resources("worker_nodes") == {}
+
+    assert config.get_node_labels("worker_nodes1") == {"foo": "bar"}
+
+    assert config.get_config("cluster_name") == "test-cli"
+    assert config.get_config("non-existing", "default") == "default"
+    assert config.get_config("non-existing") is None
+
+
+def test_multi_provider_instance_type():
+    def load_config(file):
+        path = os.path.join(AUTOSCALER_DIR_PATH, file)
+        return FileConfigReader(path).get_cached_autoscaling_config()
+
+    aws_config = load_config("aws/defaults.yaml")
+    assert aws_config.get_provider_instance_type("ray.head.default") == "m5.large"
+
+    gcp_config = load_config("gcp/defaults.yaml")
+    # NOTE: Why is this underscore....
+    assert gcp_config.get_provider_instance_type("ray_head_default") == "n1-standard-2"
+
+    aliyun_config = load_config("aliyun/defaults.yaml")
+    assert (
+        aliyun_config.get_provider_instance_type("ray.head.default") == "ecs.n4.large"
+    )
+
+    azure_config = load_config("azure/defaults.yaml")
+    assert (
+        azure_config.get_provider_instance_type("ray.head.default") == "Standard_D2s_v3"
+    )
+
+    # TODO(rickyx):
+    # We don't have kuberay and local config yet.
+
+
+def test_node_type_configs():
+    config = FileConfigReader(
+        get_test_config_path("test_ray_complex.yaml")
+    ).get_cached_autoscaling_config()
+
+    node_type_configs = config.get_node_type_configs()
+    assert config.get_max_num_worker_nodes() == 10
+    assert len(node_type_configs) == 4
+    assert node_type_configs["head_node"].max_worker_nodes == 1
+    assert node_type_configs["head_node"].min_worker_nodes == 0
+    assert node_type_configs["head_node"].resources == {}
+    assert node_type_configs["head_node"].labels == {}
+
+    assert node_type_configs["default"].max_worker_nodes == 2
+    assert node_type_configs["default"].min_worker_nodes == 0
+    assert node_type_configs["default"].resources == {}
+    assert node_type_configs["default"].labels == {}
+
+    assert node_type_configs["worker_nodes"].max_worker_nodes == 2
+    assert node_type_configs["worker_nodes"].min_worker_nodes == 1
+    assert node_type_configs["worker_nodes"].resources == {}
+    assert node_type_configs["worker_nodes"].labels == {}
+
+    assert node_type_configs["worker_nodes1"].max_worker_nodes == 2
+    assert node_type_configs["worker_nodes1"].min_worker_nodes == 1
+    assert node_type_configs["worker_nodes1"].resources == {"CPU": 2}
+    assert node_type_configs["worker_nodes1"].labels == {"foo": "bar"}
+
+
+def test_read_config():
+    # Make a temp config file from aws/defaults.yaml
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
+        # Write "aws/defaults.yaml" to the temp file
+        with open(
+            os.path.join(AUTOSCALER_DIR_PATH, "aws/defaults.yaml"), "r"
+        ) as default_file:
+            f.write(default_file.read())
+
+    config_reader = FileConfigReader(f.name)
+
+    # Check that the config is read correctly
+    assert config_reader.get_cached_autoscaling_config().provider == Provider.AWS
+
+    # Now override the file with a different provider
+    with open(f.name, "w") as f:
+        # Replace the file with "gcp/defaults.yaml"
+        with open(
+            os.path.join(AUTOSCALER_DIR_PATH, "gcp/defaults.yaml"), "r"
+        ) as default_file:
+            f.write(default_file.read())
+
+    # Still the same.
+    assert config_reader.get_cached_autoscaling_config().provider == Provider.AWS
+
+    # Reload
+    config_reader.refresh_cached_autoscaling_config()
+    assert config_reader.get_cached_autoscaling_config().provider == Provider.GCP
+
+
+def test_readonly_node_type_name_and_fallback(monkeypatch):
+    class _DummyNodeState:
+        def __init__(self, ray_node_type_name, node_id, total_resources):
+            self.ray_node_type_name = ray_node_type_name
+            self.node_id = node_id
+            self.total_resources = total_resources
+
+    class _DummyClusterState:
+        def __init__(self, node_states):
+            self.node_states = node_states
+
+    # Avoid real GCS usage.
+    monkeypatch.setattr(config_mod, "GcsClient", lambda address: object())
+    # Build a cluster with:
+    # - 1 named head type
+    # - 2 named worker types of the same type (aggregation check)
+    # - 1 worker type without name (fallback to node_id-based type)
+    unnamed_worker_id = b"\xab"
+    fallback_name = format_readonly_node_type(binary_to_hex(unnamed_worker_id))
+    nodes = [
+        _DummyNodeState(
+            "ray.head.default", b"\x01", {"CPU": 1, "node:__internal_head__": 1}
+        ),
+        _DummyNodeState("worker.custom", b"\x02", {"CPU": 2}),
+        _DummyNodeState("worker.custom", b"\x03", {"CPU": 2}),
+        _DummyNodeState("worker.custom", b"\x04", {"CPU": 2}),
+        _DummyNodeState("", unnamed_worker_id, {"CPU": 3}),
+    ]
+    monkeypatch.setattr(
+        config_mod,
+        "get_cluster_resource_state",
+        lambda _gc: _DummyClusterState(nodes),
+    )
+
+    reader = ReadOnlyProviderConfigReader("dummy:0")
+    reader.refresh_cached_autoscaling_config()
+    cfg = reader.get_cached_autoscaling_config()
+
+    node_types = cfg.get_config("available_node_types")
+    # Head assertions
+    assert "ray.head.default" in node_types
+    assert node_types["ray.head.default"]["max_workers"] == 0
+    assert cfg.get_head_node_type() == "ray.head.default"
+    # Preferred name aggregation
+    assert "worker.custom" in node_types
+    assert node_types["worker.custom"]["max_workers"] == 3
+    # Fallback for unnamed worker
+    assert fallback_name in node_types
+    assert node_types[fallback_name]["max_workers"] == 1
+
+    # Global max_workers should be the sum of all worker-type max_workers,
+    # NOT the count of node type names.
+    # Here: 3 distinct types (head, worker.custom, fallback), but
+    # 4 actual workers (3 x worker.custom + 1 x fallback).
+    # The old buggy `len(available_node_types)` returned 3 instead of 4.
+    assert cfg.get_config("max_workers") == 4
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/test_event_logger.py

@@ -0,0 +1,131 @@
+import logging
+import os
+import sys
+
+import pytest
+
+from ray.autoscaler.v2.event_logger import AutoscalerEventLogger
+from ray.autoscaler.v2.tests.util import MockEventLogger
+from ray.autoscaler.v2.utils import ResourceRequestUtil
+from ray.core.generated.autoscaler_pb2 import (
+    ClusterResourceConstraint,
+    GangResourceRequest,
+)
+from ray.core.generated.instance_manager_pb2 import LaunchRequest, TerminationRequest
+
+# coding: utf-8
+
+
+OUTDATED = TerminationRequest.Cause.OUTDATED
+IDLE = TerminationRequest.Cause.IDLE
+MAX_NUM_NODE_PER_TYPE = TerminationRequest.Cause.MAX_NUM_NODE_PER_TYPE
+MAX_NUM_NODES = TerminationRequest.Cause.MAX_NUM_NODES
+
+
+def launch_request(instance_type: str, count: int) -> LaunchRequest:
+    return LaunchRequest(
+        instance_type=instance_type,
+        count=count,
+    )
+
+
+def termination_request(
+    instance_type: str, cause: TerminationRequest.Cause
+) -> TerminationRequest:
+    return TerminationRequest(
+        instance_id="",
+        instance_type=instance_type,
+        cause=cause,
+    )
+
+
+logger = logging.getLogger(__name__)
+
+
+def test_log_scheduling_updates():
+    mock_logger = MockEventLogger(logger)
+    event_logger = AutoscalerEventLogger(mock_logger)
+
+    launch_requests = [
+        launch_request("m4.large", 2),
+        launch_request("m4.xlarge", 2),
+    ]
+    terminate_requests = [
+        termination_request("m4.large", IDLE),
+        termination_request("m4.xlarge", OUTDATED),
+    ]
+    infeasible_requests = [
+        ResourceRequestUtil.make({"CPU": 4, "GPU": 1}),
+    ] * 100 + [ResourceRequestUtil.make({"CPU": 4})]
+
+    gang_resource_requests = [
+        [
+            ResourceRequestUtil.make({"CPU": 4, "GPU": 1}),
+            ResourceRequestUtil.make({"CPU": 4, "GPU": 1}),
+        ]
+    ]
+    cluster_resource_constraints = [
+        ResourceRequestUtil.make({"CPU": 1, "GPU": 1}),
+    ] * 100
+
+    event_logger.log_cluster_scheduling_update(
+        launch_requests=launch_requests,
+        terminate_requests=terminate_requests,
+        infeasible_requests=infeasible_requests,
+        infeasible_gang_requests=[
+            GangResourceRequest(requests=reqs) for reqs in gang_resource_requests
+        ],
+        infeasible_cluster_resource_constraints=[
+            ClusterResourceConstraint(
+                resource_requests=ResourceRequestUtil.group_by_count(
+                    cluster_resource_constraints
+                )
+            )
+        ],
+        cluster_resources={"CPU": 5, "GPU": 5, "TPU": 2},
+    )
+
+    assert mock_logger.get_logs("info") == [
+        "Adding 2 node(s) of type m4.large.",
+        "Adding 2 node(s) of type m4.xlarge.",
+        "Removing 1 nodes of type m4.large (idle).",
+        "Removing 1 nodes of type m4.xlarge (outdated).",
+        "Resized to 5 CPUs, 5 GPUs, 2 TPUs.",
+    ]
+    expect_lines = [
+        "No available node types can fulfill resource requests",  # noqa
+        "No available node types can fulfill placement group requests",  # noqa
+        "No available node types can fulfill cluster constraint",  # noqa
+    ]
+    for expect_line, actual_line in zip(expect_lines, mock_logger.get_logs("error")):
+        assert expect_line in actual_line
+
+    assert mock_logger.get_logs("error") == []
+    assert mock_logger.get_logs("debug") == [
+        "Current cluster resources: {'CPU': 5, 'GPU': 5, 'TPU': 2}."
+    ]
+
+
+def test_log_scheduling_updates_without_cluster_shape():
+    mock_logger = MockEventLogger(logger)
+    event_logger = AutoscalerEventLogger(mock_logger, log_cluster_shape=False)
+
+    event_logger.log_cluster_scheduling_update(
+        launch_requests=[launch_request("m4.large", 1)],
+        terminate_requests=[termination_request("m4.xlarge", OUTDATED)],
+        infeasible_requests=[ResourceRequestUtil.make({"CPU": 4})],
+        cluster_resources={"CPU": 5},
+    )
+
+    assert mock_logger.get_logs("info") == []
+    assert mock_logger.get_logs("warning") == [
+        "No available node types can fulfill resource requests {'CPU': 4.0}*1. Add suitable node types to this cluster to resolve this issue."
+    ]
+    assert mock_logger.get_logs("debug") == []
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/test_instance_manager.py

@@ -0,0 +1,414 @@
+import os
+import sys
+import unittest
+from collections import defaultdict
+from unittest.mock import MagicMock
+
+# coding: utf-8
+import pytest
+
+from ray.autoscaler.v2.instance_manager.instance_manager import InstanceManager
+from ray.autoscaler.v2.instance_manager.instance_storage import InstanceStorage
+from ray.autoscaler.v2.instance_manager.storage import InMemoryStorage, StoreStatus
+from ray.autoscaler.v2.tests.util import MockSubscriber
+from ray.core.generated.instance_manager_pb2 import (
+    GetInstanceManagerStateRequest,
+    Instance,
+    InstanceUpdateEvent,
+    NodeKind,
+    StatusCode,
+    UpdateInstanceManagerStateRequest,
+)
+
+
+class InstanceManagerTest(unittest.TestCase):
+    def test_instances_version_mismatch(self):
+        ins_storage = MagicMock()
+        subscriber = MockSubscriber()
+        im = InstanceManager(
+            ins_storage, instance_status_update_subscribers=[subscriber]
+        )
+        # Version mismatch on reading from the storage.
+        ins_storage.get_instances.return_value = ({}, 1)
+
+        update = InstanceUpdateEvent(
+            instance_id="id-1",
+            new_instance_status=Instance.QUEUED,
+            instance_type="type-1",
+            upsert=True,
+        )
+        reply = im.update_instance_manager_state(
+            UpdateInstanceManagerStateRequest(
+                expected_version=0,
+                updates=[update],
+            )
+        )
+        assert reply.status.code == StatusCode.VERSION_MISMATCH
+        assert len(subscriber.events) == 0
+
+        # Version OK.
+        reply = im.update_instance_manager_state(
+            UpdateInstanceManagerStateRequest(
+                expected_version=1,
+                updates=[update],
+            )
+        )
+        assert reply.status.code == StatusCode.OK
+        assert len(subscriber.events) == 1
+        assert subscriber.events[0].new_instance_status == Instance.QUEUED
+
+        # Version mismatch when writing to the storage (race happens)
+        ins_storage.batch_upsert_instances.return_value = StoreStatus(
+            False, 2  # No longer 1
+        )
+        subscriber.clear()
+        reply = im.update_instance_manager_state(
+            UpdateInstanceManagerStateRequest(
+                expected_version=1,
+                updates=[update],
+            )
+        )
+        assert reply.status.code == StatusCode.VERSION_MISMATCH
+        assert len(subscriber.events) == 0
+
+        # Non-version mismatch error.
+        ins_storage.batch_upsert_instances.return_value = StoreStatus(
+            False, 1  # Still 1
+        )
+        reply = im.update_instance_manager_state(
+            UpdateInstanceManagerStateRequest(
+                expected_version=1,
+                updates=[update],
+            )
+        )
+        assert reply.status.code == StatusCode.UNKNOWN_ERRORS
+        assert len(subscriber.events) == 0
+
+    def test_get_and_updates(self):
+        ins_storage = InstanceStorage(
+            "cluster-id",
+            InMemoryStorage(),
+        )
+        subscriber = MockSubscriber()
+        im = InstanceManager(
+            ins_storage, instance_status_update_subscribers=[subscriber]
+        )
+
+        # Empty storage.
+        reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
+        assert reply.status.code == StatusCode.OK
+        assert list(reply.state.instances) == []
+
+        # Launch nodes.
+        reply = im.update_instance_manager_state(
+            UpdateInstanceManagerStateRequest(
+                expected_version=0,
+                updates=[
+                    InstanceUpdateEvent(
+                        instance_type="type-1",
+                        instance_id="id-1",
+                        new_instance_status=Instance.QUEUED,
+                        upsert=True,
+                    ),
+                    InstanceUpdateEvent(
+                        instance_type="type-2",
+                        instance_id="id-2",
+                        new_instance_status=Instance.QUEUED,
+                        upsert=True,
+                    ),
+                    InstanceUpdateEvent(
+                        instance_type="type-2",
+                        instance_id="id-3",
+                        new_instance_status=Instance.QUEUED,
+                        upsert=True,
+                    ),
+                ],
+            )
+        )
+        assert reply.status.code == StatusCode.OK
+        assert len(subscriber.events) == 3
+        for e in subscriber.events:
+            assert e.new_instance_status == Instance.QUEUED
+
+        # Get launched nodes.
+        reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
+        assert reply.status.code == StatusCode.OK
+        assert len(reply.state.instances) == 3
+
+        instance_ids = [ins.instance_id for ins in reply.state.instances]
+
+        types_count = defaultdict(int)
+        for ins in reply.state.instances:
+            types_count[ins.instance_type] += 1
+            assert ins.status == Instance.QUEUED
+
+        assert types_count["type-1"] == 1
+        assert types_count["type-2"] == 2
+
+        # Update node status.
+        subscriber.clear()
+        reply = im.update_instance_manager_state(
+            UpdateInstanceManagerStateRequest(
+                expected_version=1,
+                updates=[
+                    InstanceUpdateEvent(
+                        instance_id=instance_ids[0],
+                        new_instance_status=Instance.REQUESTED,
+                        instance_type="type-1",
+                        launch_request_id="l1",
+                    ),
+                    InstanceUpdateEvent(
+                        instance_id=instance_ids[1],
+                        new_instance_status=Instance.REQUESTED,
+                        launch_request_id="l1",
+                        instance_type="type-1",
+                    ),
+                ],
+            )
+        )
+
+        assert reply.status.code == StatusCode.OK
+        assert len(subscriber.events) == 2
+        for e in subscriber.events:
+            assert e.new_instance_status == Instance.REQUESTED
+
+        # Get updated nodes.
+        reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
+        assert reply.status.code == StatusCode.OK
+        assert len(reply.state.instances) == 3
+
+        types_count = defaultdict(int)
+        for ins in reply.state.instances:
+            types_count[ins.instance_type] += 1
+            if ins.instance_id in [instance_ids[0], instance_ids[1]]:
+                assert ins.status == Instance.REQUESTED
+            else:
+                assert ins.status == Instance.QUEUED
+
+        # Invalid instances status update.
+        subscriber.clear()
+        with pytest.raises(AssertionError):
+            reply = im.update_instance_manager_state(
+                UpdateInstanceManagerStateRequest(
+                    expected_version=2,
+                    updates=[
+                        InstanceUpdateEvent(
+                            instance_id=instance_ids[2],
+                            # Not requested yet.
+                            new_instance_status=Instance.RAY_RUNNING,
+                        ),
+                    ],
+                )
+            )
+        assert len(subscriber.events) == 0
+
+        # Invalid versions.
+        reply = im.update_instance_manager_state(
+            UpdateInstanceManagerStateRequest(
+                expected_version=0,  # Invalid version, outdated.
+                updates=[
+                    InstanceUpdateEvent(
+                        instance_id=instance_ids[2],
+                        new_instance_status=Instance.REQUESTED,
+                        instance_type="type-2",
+                    ),
+                ],
+            )
+        )
+        assert reply.status.code == StatusCode.VERSION_MISMATCH
+        assert len(subscriber.events) == 0
+
+    def test_insert(self):
+        ins_storage = InstanceStorage(
+            "cluster-id",
+            InMemoryStorage(),
+        )
+        subscriber = MockSubscriber()
+        im = InstanceManager(
+            ins_storage, instance_status_update_subscribers=[subscriber]
+        )
+
+        im.update_instance_manager_state(
+            UpdateInstanceManagerStateRequest(
+                expected_version=0,
+                updates=[
+                    InstanceUpdateEvent(
+                        instance_type="type-1",
+                        instance_id="id-1",
+                        new_instance_status=Instance.QUEUED,
+                        upsert=True,
+                    ),
+                    InstanceUpdateEvent(
+                        instance_id="id-2",
+                        new_instance_status=Instance.TERMINATING,
+                        cloud_instance_id="cloud-id-2",
+                        upsert=True,
+                    ),
+                    InstanceUpdateEvent(
+                        instance_id="id-3",
+                        new_instance_status=Instance.ALLOCATED,
+                        cloud_instance_id="cloud-id-3",
+                        node_kind=NodeKind.WORKER,
+                        instance_type="type-3",
+                        upsert=True,
+                    ),
+                ],
+            )
+        )
+        reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
+        assert len(reply.state.instances) == 3
+        instance_by_ids = {ins.instance_id: ins for ins in reply.state.instances}
+        assert instance_by_ids["id-1"].status == Instance.QUEUED
+        assert instance_by_ids["id-1"].instance_type == "type-1"
+        assert instance_by_ids["id-2"].status == Instance.TERMINATING
+        assert instance_by_ids["id-3"].status == Instance.ALLOCATED
+        assert instance_by_ids["id-3"].cloud_instance_id == "cloud-id-3"
+        version = reply.state.version
+
+        # With non-upsert flags.
+        with pytest.raises(AssertionError):
+            reply = im.update_instance_manager_state(
+                UpdateInstanceManagerStateRequest(
+                    expected_version=version,
+                    updates=[
+                        InstanceUpdateEvent(
+                            instance_type="type-1",
+                            instance_id="id-999",
+                            new_instance_status=Instance.QUEUED,
+                        ),
+                    ],
+                )
+            )
+
+        # With invalid statuses
+        all_statuses = set(Instance.InstanceStatus.values())
+        non_insertable_statuses = all_statuses - {
+            Instance.QUEUED,
+            Instance.TERMINATING,
+            Instance.ALLOCATED,
+        }
+
+        for status in non_insertable_statuses:
+            subscriber.clear()
+            with pytest.raises(AssertionError):
+                reply = im.update_instance_manager_state(
+                    UpdateInstanceManagerStateRequest(
+                        expected_version=version,
+                        updates=[
+                            InstanceUpdateEvent(
+                                instance_id="id-999",
+                                new_instance_status=status,
+                            ),
+                        ],
+                    )
+                )
+            assert len(subscriber.events) == 0
+
+    def test_apply_update(self):
+        ins_storage = InstanceStorage(
+            "cluster-id",
+            InMemoryStorage(),
+        )
+        subscriber = MockSubscriber()
+        im = InstanceManager(
+            ins_storage, instance_status_update_subscribers=[subscriber]
+        )
+
+        # Insert a new instance.
+        im.update_instance_manager_state(
+            UpdateInstanceManagerStateRequest(
+                expected_version=0,
+                updates=[
+                    InstanceUpdateEvent(
+                        instance_type="type-1",
+                        instance_id="id-1",
+                        new_instance_status=Instance.QUEUED,
+                        upsert=True,
+                    ),
+                ],
+            )
+        )
+        reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
+        assert len(reply.state.instances) == 1
+        assert reply.state.instances[0].status == Instance.QUEUED
+        assert reply.state.instances[0].instance_type == "type-1"
+
+        # Request
+        im.update_instance_manager_state(
+            UpdateInstanceManagerStateRequest(
+                expected_version=1,
+                updates=[
+                    InstanceUpdateEvent(
+                        instance_id="id-1",
+                        new_instance_status=Instance.REQUESTED,
+                        launch_request_id="l1",
+                        instance_type="type-1",
+                    ),
+                ],
+            )
+        )
+        reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
+        assert len(reply.state.instances) == 1
+        assert reply.state.instances[0].status == Instance.REQUESTED
+        assert reply.state.instances[0].launch_request_id == "l1"
+
+        # ALLOCATED
+        im.update_instance_manager_state(
+            UpdateInstanceManagerStateRequest(
+                expected_version=2,
+                updates=[
+                    InstanceUpdateEvent(
+                        instance_id="id-1",
+                        new_instance_status=Instance.ALLOCATED,
+                        cloud_instance_id="cloud-id-1",
+                        node_kind=NodeKind.WORKER,
+                        instance_type="type-1",
+                    ),
+                ],
+            )
+        )
+        reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
+        assert len(reply.state.instances) == 1
+        assert reply.state.instances[0].status == Instance.ALLOCATED
+        assert reply.state.instances[0].cloud_instance_id == "cloud-id-1"
+
+        # RAY_RUNNING
+        im.update_instance_manager_state(
+            UpdateInstanceManagerStateRequest(
+                expected_version=3,
+                updates=[
+                    InstanceUpdateEvent(
+                        instance_id="id-1",
+                        new_instance_status=Instance.RAY_RUNNING,
+                        ray_node_id="ray-node-1",
+                    ),
+                ],
+            )
+        )
+        reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
+        assert len(reply.state.instances) == 1
+        assert reply.state.instances[0].status == Instance.RAY_RUNNING
+        assert reply.state.instances[0].node_id == "ray-node-1"
+
+        # TERMINATED
+        im.update_instance_manager_state(
+            UpdateInstanceManagerStateRequest(
+                expected_version=4,
+                updates=[
+                    InstanceUpdateEvent(
+                        instance_id="id-1",
+                        new_instance_status=Instance.TERMINATED,
+                    ),
+                ],
+            )
+        )
+        reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
+        assert len(reply.state.instances) == 1
+        assert reply.state.instances[0].status == Instance.TERMINATED
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/test_instance_storage.py

@@ -0,0 +1,220 @@
+# coding: utf-8
+import copy
+import os
+import sys
+from unittest import mock
+
+import pytest  # noqa
+
+from ray.autoscaler.v2.instance_manager.instance_storage import InstanceStorage
+from ray.autoscaler.v2.instance_manager.storage import InMemoryStorage
+from ray.core.generated.instance_manager_pb2 import Instance
+
+
+def create_instance(instance_id, status=Instance.QUEUED, version=0):
+    return Instance(
+        instance_id=instance_id,
+        status=status,
+        version=version,
+    )
+
+
+@mock.patch("time.time", mock.MagicMock(return_value=1))
+def test_upsert():
+    storage = InstanceStorage(
+        cluster_id="test_cluster",
+        storage=InMemoryStorage(),
+    )
+    instance1 = create_instance("instance1")
+    instance2 = create_instance("instance2")
+    instance3 = create_instance("instance3")
+
+    assert (True, 1) == storage.batch_upsert_instances(
+        [instance1, instance2],
+        expected_storage_version=None,
+    )
+
+    instance1.version = 1
+    instance2.version = 1
+    entries, storage_version = storage.get_instances()
+
+    assert storage_version == 1
+    assert entries == {
+        "instance1": instance1,
+        "instance2": instance2,
+    }
+
+    assert (False, 1) == storage.batch_upsert_instances(
+        [create_instance("instance1"), create_instance("instance2")],
+        expected_storage_version=0,
+    )
+
+    instance2.status = Instance.ALLOCATED
+    assert (True, 2) == storage.batch_upsert_instances(
+        [instance3, instance2],
+        expected_storage_version=1,
+    )
+
+    instance1.version = 1
+    instance2.version = 2
+    instance3.version = 2
+    entries, storage_version = storage.get_instances()
+
+    assert storage_version == 2
+    assert entries == {
+        "instance1": instance1,
+        "instance2": instance2,
+        "instance3": instance3,
+    }
+
+
+@mock.patch("time.time", mock.MagicMock(return_value=1))
+def test_update():
+
+    storage = InstanceStorage(
+        cluster_id="test_cluster",
+        storage=InMemoryStorage(),
+    )
+    instance1 = create_instance("instance1")
+    instance2 = create_instance("instance2")
+
+    assert (True, 1) == storage.upsert_instance(instance=instance1)
+    assert (True, 2) == storage.upsert_instance(instance=instance2)
+
+    assert (
+        {
+            "instance1": create_instance("instance1", version=1),
+            "instance2": create_instance("instance2", version=2),
+        },
+        2,
+    ) == storage.get_instances()
+
+    # failed because instance version is not correct
+    assert (False, 2) == storage.upsert_instance(
+        instance=instance1,
+        expected_instance_version=0,
+    )
+
+    # failed because storage version is not correct
+    assert (False, 2) == storage.upsert_instance(
+        instance=instance1,
+        expected_storage_version=0,
+    )
+
+    assert (True, 3) == storage.upsert_instance(
+        instance=instance2,
+        expected_storage_version=2,
+    )
+
+    assert (
+        {
+            "instance1": create_instance("instance1", version=1),
+            "instance2": create_instance("instance2", version=3),
+        },
+        3,
+    ) == storage.get_instances()
+
+    assert (True, 4) == storage.upsert_instance(
+        instance=instance1,
+        expected_instance_version=1,
+    )
+
+    assert (
+        {
+            "instance1": create_instance("instance1", version=4),
+            "instance2": create_instance("instance2", version=3),
+        },
+        4,
+    ) == storage.get_instances()
+
+
+@mock.patch("time.time", mock.MagicMock(return_value=1))
+def test_delete():
+
+    storage = InstanceStorage(
+        cluster_id="test_cluster",
+        storage=InMemoryStorage(),
+    )
+    instance1 = create_instance("instance1")
+    instance2 = create_instance("instance2")
+    instance3 = create_instance("instance3")
+
+    assert (True, 1) == storage.batch_upsert_instances(
+        [instance1, instance2, instance3],
+        expected_storage_version=None,
+    )
+
+    assert (False, 1) == storage.batch_delete_instances(
+        instance_ids=["instance1"], expected_storage_version=0
+    )
+    assert (True, 2) == storage.batch_delete_instances(instance_ids=["instance1"])
+
+    assert (
+        {
+            "instance2": create_instance("instance2", version=1),
+            "instance3": create_instance("instance3", version=1),
+        },
+        2,
+    ) == storage.get_instances()
+
+    assert (True, 3) == storage.batch_delete_instances(
+        instance_ids=["instance2"], expected_storage_version=2
+    )
+
+    assert (
+        {
+            "instance3": create_instance("instance3", version=1),
+        },
+        3,
+    ) == storage.get_instances()
+
+
+@mock.patch("time.time", mock.MagicMock(return_value=1))
+def test_get_instances():
+    storage = InstanceStorage(
+        cluster_id="test_cluster",
+        storage=InMemoryStorage(),
+    )
+    instance1 = create_instance("instance1", version=1)
+    instance2 = create_instance("instance2", status=Instance.ALLOCATED, version=1)
+    instance3 = create_instance("instance3", status=Instance.TERMINATING, version=1)
+
+    assert (True, 1) == storage.batch_upsert_instances(
+        [copy.deepcopy(instance1), copy.deepcopy(instance2), copy.deepcopy(instance3)],
+        expected_storage_version=None,
+    )
+
+    assert (
+        {
+            "instance1": instance1,
+            "instance2": instance2,
+            "instance3": instance3,
+        },
+        1,
+    ) == storage.get_instances()
+
+    assert (
+        {
+            "instance1": instance1,
+            "instance2": instance2,
+        },
+        1,
+    ) == storage.get_instances(instance_ids=["instance1", "instance2"])
+
+    assert ({"instance2": instance2}, 1) == storage.get_instances(
+        instance_ids=["instance1", "instance2"], status_filter={Instance.ALLOCATED}
+    )
+
+    assert (
+        {
+            "instance2": instance2,
+        },
+        1,
+    ) == storage.get_instances(status_filter={Instance.ALLOCATED})
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/test_instance_util.py

@@ -0,0 +1,336 @@
+import os
+import sys
+import unittest
+
+# coding: utf-8
+from unittest.mock import patch
+
+import pytest
+
+from ray.autoscaler.v2.instance_manager.common import InstanceUtil
+from ray.core.generated.instance_manager_pb2 import Instance
+
+
+class InstanceUtilTest(unittest.TestCase):
+    def test_basic(self):
+        # New instance.
+        instance = InstanceUtil.new_instance("i-123", "type_1", Instance.QUEUED)
+        assert instance.instance_id == "i-123"
+        assert instance.instance_type == "type_1"
+        assert instance.status == Instance.QUEUED
+
+        # Set status.
+        assert InstanceUtil.set_status(instance, Instance.REQUESTED)
+        assert instance.status == Instance.REQUESTED
+
+        # Set status with invalid status.
+        assert not InstanceUtil.set_status(instance, Instance.RAY_RUNNING)
+
+        assert not InstanceUtil.set_status(instance, Instance.UNKNOWN)
+
+    def test_transition_graph(self):
+        # Assert on each edge in the graph.
+        all_status = set(Instance.InstanceStatus.values())
+
+        g = InstanceUtil.get_valid_transitions()
+
+        assert g[Instance.QUEUED] == {Instance.REQUESTED, Instance.TERMINATED}
+        all_status.remove(Instance.QUEUED)
+
+        assert g[Instance.REQUESTED] == {
+            Instance.ALLOCATED,
+            Instance.QUEUED,
+            Instance.ALLOCATION_FAILED,
+        }
+        all_status.remove(Instance.REQUESTED)
+
+        assert g[Instance.ALLOCATED] == {
+            Instance.RAY_INSTALLING,
+            Instance.RAY_RUNNING,
+            Instance.RAY_STOPPING,
+            Instance.RAY_STOPPED,
+            Instance.TERMINATING,
+            Instance.TERMINATED,
+            Instance.ALLOCATION_TIMEOUT,
+        }
+        all_status.remove(Instance.ALLOCATED)
+
+        assert g[Instance.RAY_INSTALLING] == {
+            Instance.RAY_RUNNING,
+            Instance.RAY_INSTALL_FAILED,
+            Instance.RAY_STOPPED,
+            Instance.TERMINATING,
+            Instance.TERMINATED,
+        }
+        all_status.remove(Instance.RAY_INSTALLING)
+
+        assert g[Instance.RAY_RUNNING] == {
+            Instance.RAY_STOP_REQUESTED,
+            Instance.RAY_STOPPING,
+            Instance.RAY_STOPPED,
+            Instance.TERMINATING,
+            Instance.TERMINATED,
+        }
+        all_status.remove(Instance.RAY_RUNNING)
+
+        assert g[Instance.ALLOCATION_TIMEOUT] == {
+            Instance.TERMINATING,
+            Instance.TERMINATED,
+        }
+        all_status.remove(Instance.ALLOCATION_TIMEOUT)
+
+        assert g[Instance.RAY_STOP_REQUESTED] == {
+            Instance.RAY_STOPPING,
+            Instance.RAY_STOPPED,
+            Instance.TERMINATED,
+            Instance.RAY_RUNNING,
+        }
+        all_status.remove(Instance.RAY_STOP_REQUESTED)
+
+        assert g[Instance.RAY_STOPPING] == {
+            Instance.RAY_STOPPED,
+            Instance.TERMINATING,
+            Instance.TERMINATED,
+        }
+
+        all_status.remove(Instance.RAY_STOPPING)
+
+        assert g[Instance.RAY_STOPPED] == {Instance.TERMINATED, Instance.TERMINATING}
+        all_status.remove(Instance.RAY_STOPPED)
+
+        assert g[Instance.TERMINATING] == {
+            Instance.TERMINATED,
+            Instance.TERMINATION_FAILED,
+        }
+        all_status.remove(Instance.TERMINATING)
+
+        assert g[Instance.TERMINATION_FAILED] == {
+            Instance.TERMINATING,
+            Instance.TERMINATED,
+        }
+        all_status.remove(Instance.TERMINATION_FAILED)
+
+        assert g[Instance.TERMINATED] == set()
+        all_status.remove(Instance.TERMINATED)
+
+        assert g[Instance.ALLOCATION_FAILED] == set()
+        all_status.remove(Instance.ALLOCATION_FAILED)
+
+        assert g[Instance.RAY_INSTALL_FAILED] == {
+            Instance.TERMINATED,
+            Instance.TERMINATING,
+        }
+        all_status.remove(Instance.RAY_INSTALL_FAILED)
+
+        assert g[Instance.UNKNOWN] == set()
+        all_status.remove(Instance.UNKNOWN)
+
+        assert len(all_status) == 0
+
+    @patch("time.time_ns")
+    def test_status_time(self, mock_time):
+        mock_time.return_value = 1
+        instance = InstanceUtil.new_instance("i-123", "type_1", Instance.QUEUED)
+        # OK
+        assert (
+            InstanceUtil.get_status_transition_times_ns(instance, Instance.QUEUED)[0]
+            == 1
+        )
+        # No filter.
+        assert InstanceUtil.get_status_transition_times_ns(
+            instance,
+        ) == [1]
+
+        # Missing status returns empty list
+        assert (
+            InstanceUtil.get_status_transition_times_ns(instance, Instance.REQUESTED)
+            == []
+        )
+
+        # Multiple status.
+        mock_time.return_value = 2
+        InstanceUtil.set_status(instance, Instance.REQUESTED)
+        mock_time.return_value = 3
+        InstanceUtil.set_status(instance, Instance.QUEUED)
+        mock_time.return_value = 4
+        InstanceUtil.set_status(instance, Instance.REQUESTED)
+        assert InstanceUtil.get_status_transition_times_ns(
+            instance, Instance.QUEUED
+        ) == [1, 3]
+
+    @patch("time.time_ns")
+    def test_get_last_status_transition(self, mock_time):
+        mock_time.return_value = 1
+        instance = InstanceUtil.new_instance("i-123", "type_1", Instance.QUEUED)
+        assert (
+            InstanceUtil.get_last_status_transition(instance).instance_status
+            == Instance.QUEUED
+        )
+        assert InstanceUtil.get_last_status_transition(instance).timestamp_ns == 1
+
+        mock_time.return_value = 2
+        InstanceUtil.set_status(instance, Instance.REQUESTED)
+        assert (
+            InstanceUtil.get_last_status_transition(instance).instance_status
+            == Instance.REQUESTED
+        )
+        assert InstanceUtil.get_last_status_transition(instance).timestamp_ns == 2
+
+        mock_time.return_value = 3
+        InstanceUtil.set_status(instance, Instance.QUEUED)
+        assert (
+            InstanceUtil.get_last_status_transition(instance).instance_status
+            == Instance.QUEUED
+        )
+        assert InstanceUtil.get_last_status_transition(instance).timestamp_ns == 3
+
+        assert (
+            InstanceUtil.get_last_status_transition(
+                instance, select_instance_status=Instance.REQUESTED
+            ).instance_status
+            == Instance.REQUESTED
+        )
+        assert (
+            InstanceUtil.get_last_status_transition(
+                instance, select_instance_status=Instance.REQUESTED
+            ).timestamp_ns
+            == 2
+        )
+
+        assert (
+            InstanceUtil.get_last_status_transition(
+                instance, select_instance_status=Instance.RAY_RUNNING
+            )
+            is None
+        )
+
+    def test_is_cloud_instance_allocated(self):
+        all_status = set(Instance.InstanceStatus.values())
+        instance = InstanceUtil.new_instance("i-123", "type_1", Instance.QUEUED)
+        positive_status = {
+            Instance.ALLOCATED,
+            Instance.RAY_INSTALLING,
+            Instance.RAY_INSTALL_FAILED,
+            Instance.RAY_RUNNING,
+            Instance.RAY_STOP_REQUESTED,
+            Instance.RAY_STOPPING,
+            Instance.RAY_STOPPED,
+            Instance.TERMINATING,
+            Instance.TERMINATION_FAILED,
+            Instance.ALLOCATION_TIMEOUT,
+        }
+        for s in positive_status:
+            instance.status = s
+            assert InstanceUtil.is_cloud_instance_allocated(instance.status)
+            all_status.remove(s)
+
+        # Unknown not possible.
+        all_status.remove(Instance.UNKNOWN)
+        for s in all_status:
+            instance.status = s
+            assert not InstanceUtil.is_cloud_instance_allocated(instance.status)
+
+    def test_is_ray_running(self):
+        all_statuses = set(Instance.InstanceStatus.values())
+        positive_statuses = {
+            Instance.RAY_RUNNING,
+            Instance.RAY_STOP_REQUESTED,
+            Instance.RAY_STOPPING,
+        }
+        all_statuses.remove(Instance.UNKNOWN)
+        for s in positive_statuses:
+            assert InstanceUtil.is_ray_running(s)
+            all_statuses.remove(s)
+
+        for s in all_statuses:
+            assert not InstanceUtil.is_ray_running(s)
+
+    def test_is_ray_pending(self):
+        all_statuses = set(Instance.InstanceStatus.values())
+        all_statuses.remove(Instance.UNKNOWN)
+        positive_statuses = {
+            Instance.QUEUED,
+            Instance.REQUESTED,
+            Instance.RAY_INSTALLING,
+            Instance.ALLOCATED,
+        }
+        for s in positive_statuses:
+            assert InstanceUtil.is_ray_pending(s), Instance.InstanceStatus.Name(s)
+            all_statuses.remove(s)
+
+        for s in all_statuses:
+            assert not InstanceUtil.is_ray_pending(s), Instance.InstanceStatus.Name(s)
+
+    def test_is_ray_running_reachable(self):
+        all_status = set(Instance.InstanceStatus.values())
+        positive_status = {
+            Instance.QUEUED,
+            Instance.REQUESTED,
+            Instance.ALLOCATED,
+            Instance.RAY_INSTALLING,
+            Instance.RAY_RUNNING,
+            Instance.RAY_STOP_REQUESTED,
+        }
+        for s in positive_status:
+            assert InstanceUtil.is_ray_running_reachable(
+                s
+            ), Instance.InstanceStatus.Name(s)
+            all_status.remove(s)
+
+        # Unknown not possible.
+        all_status.remove(Instance.UNKNOWN)
+        for s in all_status:
+            assert not InstanceUtil.is_ray_running_reachable(
+                s
+            ), Instance.InstanceStatus.Name(s)
+
+    def test_reachable_from(self):
+        def add_reachable_from(reachable, src, transitions):
+            reachable[src] = set()
+            for dst in transitions[src]:
+                reachable[src].add(dst)
+                reachable[src] |= (
+                    reachable[dst] if reachable[dst] is not None else set()
+                )
+
+        expected_reachable = {s: None for s in Instance.InstanceStatus.values()}
+
+        # Error status and terminal status.
+        expected_reachable[Instance.ALLOCATION_FAILED] = set()
+        expected_reachable[Instance.UNKNOWN] = set()
+        expected_reachable[Instance.TERMINATED] = set()
+
+        transitions = InstanceUtil.get_valid_transitions()
+
+        # Recursively build the reachable set from terminal statuses.
+        add_reachable_from(expected_reachable, Instance.TERMINATION_FAILED, transitions)
+        add_reachable_from(expected_reachable, Instance.TERMINATING, transitions)
+        # Add TERMINATION_FAILED again since it's also reachable from TERMINATING.
+        add_reachable_from(expected_reachable, Instance.TERMINATION_FAILED, transitions)
+        add_reachable_from(expected_reachable, Instance.RAY_STOPPED, transitions)
+        add_reachable_from(expected_reachable, Instance.RAY_STOPPING, transitions)
+        add_reachable_from(expected_reachable, Instance.RAY_STOP_REQUESTED, transitions)
+        add_reachable_from(expected_reachable, Instance.RAY_RUNNING, transitions)
+        # Add RAY_STOP_REQUESTED again since it's also reachable from RAY_RUNNING.
+        add_reachable_from(expected_reachable, Instance.RAY_STOP_REQUESTED, transitions)
+        add_reachable_from(expected_reachable, Instance.RAY_INSTALL_FAILED, transitions)
+        add_reachable_from(expected_reachable, Instance.RAY_INSTALLING, transitions)
+        add_reachable_from(expected_reachable, Instance.ALLOCATED, transitions)
+        add_reachable_from(expected_reachable, Instance.REQUESTED, transitions)
+        add_reachable_from(expected_reachable, Instance.QUEUED, transitions)
+        # Add REQUESTED again since it's also reachable from QUEUED.
+        add_reachable_from(expected_reachable, Instance.REQUESTED, transitions)
+        add_reachable_from(expected_reachable, Instance.ALLOCATION_TIMEOUT, transitions)
+
+        for s, expected in expected_reachable.items():
+            assert InstanceUtil.get_reachable_statuses(s) == expected, (
+                f"reachable_from({s}) = {InstanceUtil.get_reachable_statuses(s)} "
+                f"!= {expected}"
+            )
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/test_metrics_reporter.py

@@ -0,0 +1,219 @@
+import os
+import sys
+from typing import List
+
+import pytest
+
+from ray.autoscaler._private.prom_metrics import AutoscalerPrometheusMetrics
+from ray.autoscaler.v2.instance_manager.config import NodeTypeConfig
+from ray.autoscaler.v2.metrics_reporter import AutoscalerMetricsReporter
+from ray.autoscaler.v2.tests.util import create_instance
+from ray.core.generated.instance_manager_pb2 import Instance
+
+
+def _get_metrics(metrics, name) -> List[float]:
+    sample_values = []
+    for x in metrics:
+        for sample in x.samples:
+            if sample.name == name:
+                sample_values.append(sample.value)
+    return sample_values
+
+
+def test_report_nodes_resources():
+    """
+    Test that the metrics reporter reports the correct number of nodes and resources
+    """
+    reporter = AutoscalerMetricsReporter(
+        AutoscalerPrometheusMetrics(session_name="test")
+    )
+    node_type_configs = {
+        "type_1": NodeTypeConfig(
+            name="type_1",
+            max_worker_nodes=10,
+            min_worker_nodes=1,
+            resources={"CPU": 1},
+        ),
+        "type_2": NodeTypeConfig(
+            name="type_2",
+            max_worker_nodes=10,
+            min_worker_nodes=1,
+            resources={"GPU": 1},
+        ),
+    }
+
+    _i = 0
+
+    def id():
+        nonlocal _i
+        _i += 1
+        return f"i-{_i}"
+
+    terminating_type_1 = create_instance(
+        id(), status=Instance.TERMINATING, instance_type="type_1"
+    )
+    terminating_type_2 = create_instance(
+        id(), status=Instance.TERMINATING, instance_type="type_2"
+    )
+
+    instances = [
+        # Active = 3
+        create_instance(id(), status=Instance.RAY_RUNNING, instance_type="type_1"),
+        create_instance(
+            id(), status=Instance.RAY_STOP_REQUESTED, instance_type="type_1"
+        ),
+        create_instance(id(), status=Instance.RAY_STOPPING, instance_type="type_1"),
+        create_instance(id(), status=Instance.RAY_RUNNING, instance_type="type_2"),
+        # Pending
+        create_instance(id(), status=Instance.QUEUED, instance_type="type_1"),
+        create_instance(id(), status=Instance.REQUESTED, instance_type="type_1"),
+        create_instance(id(), status=Instance.RAY_INSTALLING, instance_type="type_1"),
+        create_instance(id(), status=Instance.ALLOCATED, instance_type="type_1"),
+        create_instance(id(), status=Instance.RAY_INSTALLING, instance_type="type_2"),
+        create_instance(id(), status=Instance.ALLOCATED, instance_type="type_2"),
+        # Terminating
+        terminating_type_1,
+        terminating_type_2,
+    ]
+    reporter.report_instances(instances, node_type_configs)
+
+    assert _get_metrics(
+        reporter._prom_metrics.active_nodes.labels(
+            SessionName="test", NodeType="type_1"
+        ).collect(),
+        "autoscaler_active_nodes",
+    ) == [3]
+    assert _get_metrics(
+        reporter._prom_metrics.pending_nodes.labels(
+            SessionName="test", NodeType="type_1"
+        ).collect(),
+        "autoscaler_pending_nodes",
+    ) == [4]
+
+    assert _get_metrics(
+        reporter._prom_metrics.recently_failed_nodes.labels(
+            SessionName="test", NodeType="type_1"
+        ).collect(),
+        "autoscaler_recently_failed_nodes",
+    ) == [1]
+
+    # Test that resources are reported correctly
+    reporter.report_resources(instances, node_type_configs)
+
+    assert _get_metrics(
+        reporter._prom_metrics.cluster_resources.labels(
+            SessionName="test", resource="CPU"
+        ).collect(),
+        "autoscaler_cluster_resources",
+    ) == [3]
+
+    assert _get_metrics(
+        reporter._prom_metrics.cluster_resources.labels(
+            SessionName="test", resource="GPU"
+        ).collect(),
+        "autoscaler_cluster_resources",
+    ) == [1]
+
+    assert _get_metrics(
+        reporter._prom_metrics.pending_resources.labels(
+            SessionName="test", resource="CPU"
+        ).collect(),
+        "autoscaler_pending_resources",
+    ) == [4]
+
+    assert _get_metrics(
+        reporter._prom_metrics.pending_resources.labels(
+            SessionName="test", resource="GPU"
+        ).collect(),
+        "autoscaler_pending_resources",
+    ) == [2]
+
+
+def test_report_nodes_resources_handles_deleted_node_type():
+    reporter = AutoscalerMetricsReporter(
+        AutoscalerPrometheusMetrics(session_name="test_deleted")
+    )
+    node_type_configs = {
+        "current_type": NodeTypeConfig(
+            name="current_type",
+            max_worker_nodes=10,
+            min_worker_nodes=1,
+            resources={"CPU": 2},
+        ),
+    }
+
+    instances = [
+        create_instance(
+            "current-running",
+            status=Instance.RAY_RUNNING,
+            instance_type="current_type",
+        ),
+        create_instance(
+            "deleted-terminated",
+            status=Instance.TERMINATED,
+            instance_type="deleted_type",
+        ),
+        create_instance(
+            "deleted-terminating",
+            status=Instance.TERMINATING,
+            instance_type="deleted_type",
+        ),
+        create_instance(
+            "deleted-pending",
+            status=Instance.QUEUED,
+            instance_type="deleted_type",
+        ),
+    ]
+
+    reporter.report_instances(instances, node_type_configs)
+    reporter.report_resources(instances, node_type_configs)
+
+    assert _get_metrics(
+        reporter._prom_metrics.active_nodes.labels(
+            SessionName="test_deleted", NodeType="current_type"
+        ).collect(),
+        "autoscaler_active_nodes",
+    ) == [1]
+    assert _get_metrics(
+        reporter._prom_metrics.pending_nodes.labels(
+            SessionName="test_deleted", NodeType="current_type"
+        ).collect(),
+        "autoscaler_pending_nodes",
+    ) == [0]
+    assert _get_metrics(
+        reporter._prom_metrics.recently_failed_nodes.labels(
+            SessionName="test_deleted", NodeType="current_type"
+        ).collect(),
+        "autoscaler_recently_failed_nodes",
+    ) == [0]
+    assert _get_metrics(
+        reporter._prom_metrics.active_nodes.labels(
+            SessionName="test_deleted", NodeType="deleted_type"
+        ).collect(),
+        "autoscaler_active_nodes",
+    ) == [0]
+    assert _get_metrics(
+        reporter._prom_metrics.pending_nodes.labels(
+            SessionName="test_deleted", NodeType="deleted_type"
+        ).collect(),
+        "autoscaler_pending_nodes",
+    ) == [1]
+    assert _get_metrics(
+        reporter._prom_metrics.recently_failed_nodes.labels(
+            SessionName="test_deleted", NodeType="deleted_type"
+        ).collect(),
+        "autoscaler_recently_failed_nodes",
+    ) == [1]
+    assert _get_metrics(
+        reporter._prom_metrics.cluster_resources.labels(
+            SessionName="test_deleted", resource="CPU"
+        ).collect(),
+        "autoscaler_cluster_resources",
+    ) == [2]
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/test_priority_selection.py

@@ -0,0 +1,328 @@
+import os
+import sys
+import time
+
+import pytest
+
+from ray.autoscaler.v2.instance_manager.config import NodeTypeConfig
+from ray.autoscaler.v2.instance_manager.subscribers.cloud_resource_monitor import (
+    CloudResourceMonitor,
+)
+from ray.autoscaler.v2.scheduler import (
+    ResourceDemandScheduler,
+    ResourceRequestSource,
+    SchedulingNode,
+    SchedulingNodeStatus,
+    SchedulingRequest,
+)
+from ray.core.generated.autoscaler_pb2 import ResourceRequest as PBResourceRequest
+from ray.core.generated.instance_manager_pb2 import (
+    Instance,
+    InstanceUpdateEvent,
+    NodeKind,
+)
+
+
+def test_recovery_scoring():
+    monitor = CloudResourceMonitor()
+    node_type = "gpu-node"
+
+    # Mock failure
+    event = InstanceUpdateEvent(
+        instance_type=node_type, new_instance_status=Instance.ALLOCATION_TIMEOUT
+    )
+    monitor.notify([event])
+
+    # Immediately after failure, score should be 0.0
+    scores = monitor.get_recoverable_resource_availabilities()
+    assert scores[node_type] == 0.0
+
+    # After safety floor (e.g., 11s, default safety floor is 10s)
+    monitor._last_unavailable_timestamp[node_type] = time.time() - 11
+    scores = monitor.get_recoverable_resource_availabilities()
+    assert 0.0 < scores[node_type] < 0.1
+
+    # Halfway through recovery window (600s / 2 = 300s)
+    monitor._last_unavailable_timestamp[node_type] = time.time() - 300
+    scores = monitor.get_recoverable_resource_availabilities()
+    # 300 / 600 = 0.5
+    assert pytest.approx(scores[node_type], 0.01) == 0.5
+
+    # After recovery window
+    monitor._last_unavailable_timestamp[node_type] = time.time() - 601
+    scores = monitor.get_recoverable_resource_availabilities()
+    assert scores[node_type] == 1.0
+
+
+def test_scheduler_priority_tie_breaking():
+    # Two node types with identical resources
+    resources = {"CPU": 4}
+    node_type_1 = "high-priority"
+    node_type_2 = "low-priority"
+
+    config_1 = NodeTypeConfig(
+        name=node_type_1,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=10,
+    )
+    config_2 = NodeTypeConfig(
+        name=node_type_2,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=0,
+    )
+
+    node_1 = SchedulingNode.from_node_config(
+        config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+    node_2 = SchedulingNode.from_node_config(
+        config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+
+    request = PBResourceRequest(resources_bundle={"CPU": 1})
+
+    # Utilization and Availability are equal (all perfect)
+    # Priority should break the tie
+    best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
+        [request], [node_2, node_1], ResourceRequestSource.PENDING_DEMAND, {}, {}
+    )
+
+    assert best_node.node_type == node_type_1
+
+
+def test_schedule_context_propagation():
+    resources = {"CPU": 4}
+    node_type = "gpu-node"
+
+    config = NodeTypeConfig(
+        name=node_type,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=7,
+    )
+
+    # Mock cloud availabilities
+    cloud_availabilities = {node_type: 0.1}  # failure recency
+    recoverable_availabilities = {node_type: 0.5}
+
+    req = SchedulingRequest(
+        node_type_configs={node_type: config},
+        cloud_resource_availabilities=cloud_availabilities,
+        recoverable_resource_availabilities=recoverable_availabilities,
+        disable_launch_config_check=True,
+    )
+
+    ctx = ResourceDemandScheduler.ScheduleContext.from_schedule_request(req)
+
+    # Check if a new node created from this context has the correct priority
+    node_pools = [
+        SchedulingNode.from_node_config(
+            ctx.get_node_type_configs()[nt],
+            status=SchedulingNodeStatus.TO_LAUNCH,
+            node_kind=NodeKind.WORKER,
+        )
+        for nt, num_available in ctx.get_node_type_available().items()
+    ]
+
+    assert len(node_pools) == 1
+    node = node_pools[0]
+    assert node.priority == 7
+    # Dynamic scores are in the context, not the node
+    assert ctx.get_recoverable_resource_availabilities()[node_type] == 0.5
+    assert ctx.get_cloud_resource_availabilities()[node_type] == 0.1
+
+
+def test_scheduler_availability_over_priority():
+    # High priority node is recovering (score 0.5)
+    # Low priority node is available (score 1.0)
+    resources = {"CPU": 4}
+    node_type_1 = "high-priority"
+    node_type_2 = "low-priority"
+
+    config_1 = NodeTypeConfig(
+        name=node_type_1,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=10,
+    )
+    config_2 = NodeTypeConfig(
+        name=node_type_2,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=0,
+    )
+
+    node_1 = SchedulingNode.from_node_config(
+        config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+    node_2 = SchedulingNode.from_node_config(
+        config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+
+    request = PBResourceRequest(resources_bundle={"CPU": 1})
+
+    # Recoverable Availability is higher for node_2, so it should be chosen despite lower priority
+    best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
+        [request],
+        [node_1, node_2],
+        ResourceRequestSource.PENDING_DEMAND,
+        cloud_resource_availabilities={},
+        recoverable_resource_availabilities={node_type_1: 0.5, node_type_2: 1.0},
+    )
+
+    assert best_node.node_type == node_type_2
+
+
+def test_scheduler_failure_recency_tie_breaking():
+    # Same priority, same recoverable availability (1.0)
+    # One has an older failure than the other.
+    resources = {"CPU": 4}
+    node_type_1 = "older-failure"
+    node_type_2 = "newer-failure"
+
+    config_1 = NodeTypeConfig(
+        name=node_type_1,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=5,
+    )
+    config_2 = NodeTypeConfig(
+        name=node_type_2,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=5,
+    )
+
+    node_1 = SchedulingNode.from_node_config(
+        config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+    node_2 = SchedulingNode.from_node_config(
+        config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+
+    request = PBResourceRequest(resources_bundle={"CPU": 1})
+
+    best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
+        [request],
+        [node_1, node_2],
+        ResourceRequestSource.PENDING_DEMAND,
+        cloud_resource_availabilities={node_type_1: 0.9, node_type_2: 0.1},
+        recoverable_resource_availabilities={node_type_1: 1.0, node_type_2: 1.0},
+    )
+
+    assert best_node.node_type == node_type_1
+
+
+def test_recovery_integration():
+    monitor = CloudResourceMonitor()
+    node_type_1 = "high-priority"
+    node_type_2 = "low-priority"
+
+    # Mock failure for node_1 (high priority)
+    event = InstanceUpdateEvent(
+        instance_type=node_type_1, new_instance_status=Instance.ALLOCATION_TIMEOUT
+    )
+    monitor.notify([event])
+
+    # Score should be 0.0 for node_1 immediately
+    scores = monitor.get_recoverable_resource_availabilities()
+    assert scores[node_type_1] == 0.0
+
+    # Setup scheduler structures
+    resources = {"CPU": 4}
+    config_1 = NodeTypeConfig(
+        name=node_type_1,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=10,
+    )
+    config_2 = NodeTypeConfig(
+        name=node_type_2,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources,
+        priority=0,
+    )
+
+    node_1 = SchedulingNode.from_node_config(
+        config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+    node_2 = SchedulingNode.from_node_config(
+        config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+
+    request = PBResourceRequest(resources_bundle={"CPU": 1})
+
+    # Pass scores from monitor to scheduler
+    best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
+        [request],
+        [node_1, node_2],
+        ResourceRequestSource.PENDING_DEMAND,
+        cloud_resource_availabilities={},
+        recoverable_resource_availabilities=scores,
+    )
+
+    # Node 2 should be chosen because Node 1 is recovering, despite Node 1 having higher priority
+    assert best_node.node_type == node_type_2
+
+
+def test_scheduler_utilization_over_priority():
+    # Node 1: 4 CPUs, priority 10
+    # Node 2: 2 CPUs, priority 0
+    # Request: 2 CPUs
+    # Node 2 should be selected because it fits perfectly (utilization score is higher),
+    # even though Node 1 has higher priority.
+    resources_1 = {"CPU": 4}
+    resources_2 = {"CPU": 2}
+    node_type_1 = "large-node"
+    node_type_2 = "small-node"
+
+    config_1 = NodeTypeConfig(
+        name=node_type_1,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources_1,
+        priority=10,
+    )
+    config_2 = NodeTypeConfig(
+        name=node_type_2,
+        min_worker_nodes=0,
+        max_worker_nodes=10,
+        resources=resources_2,
+        priority=0,
+    )
+
+    node_1 = SchedulingNode.from_node_config(
+        config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+    node_2 = SchedulingNode.from_node_config(
+        config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
+    )
+
+    request = PBResourceRequest(resources_bundle={"CPU": 2})
+
+    best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
+        [request],
+        [node_1, node_2],
+        ResourceRequestSource.PENDING_DEMAND,
+        cloud_resource_availabilities={},
+        recoverable_resource_availabilities={},
+    )
+
+    assert best_node.node_type == node_type_2
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/test_ray_installer.py

@@ -0,0 +1,68 @@
+# coding: utf-8
+import os
+import sys
+import unittest
+
+import pytest  # noqa
+
+from ray._private.test_utils import load_test_config
+from ray.autoscaler.tags import TAG_RAY_NODE_KIND
+from ray.autoscaler.v2.instance_manager.config import AutoscalingConfig
+from ray.autoscaler.v2.instance_manager.ray_installer import RayInstaller
+from ray.core.generated.instance_manager_pb2 import Instance
+from ray.tests.autoscaler_test_utils import MockProcessRunner, MockProvider
+
+
+class RayInstallerTest(unittest.TestCase):
+    def setUp(self):
+        self.base_provider = MockProvider()
+        self.config = AutoscalingConfig(load_test_config("test_ray_complex.yaml"))
+        self.runner = MockProcessRunner()
+        self.ray_installer = RayInstaller(self.base_provider, self.config, self.runner)
+
+    def test_install_succeeded(self):
+        self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
+        self.runner.respond_to_call("json .Config.Env", ["[]" for i in range(1)])
+
+        self.ray_installer.install_ray(
+            Instance(
+                instance_id="0", instance_type="worker_nodes1", cloud_instance_id="0"
+            ),
+            head_node_ip="1.2.3.4",
+        )
+
+    def test_install_failed(self):
+        # creation failed because no such node.
+        with self.assertRaisesRegex(KeyError, "0"):
+            assert not self.ray_installer.install_ray(
+                Instance(
+                    instance_id="0",
+                    instance_type="worker_nodes1",
+                    cloud_instance_id="0",
+                ),
+                head_node_ip="1.2.3.4",
+            )
+
+        self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
+        self.runner.fail_cmds = [
+            "echo"  # this is the command used in the test_ray_complex.yaml
+        ]
+        self.runner.respond_to_call("json .Config.Env", ["[]" for i in range(1)])
+
+        # creation failed because setup command failed.
+        with self.assertRaisesRegex(Exception, "unexpected status"):
+            self.ray_installer.install_ray(
+                Instance(
+                    instance_id="0",
+                    instance_type="worker_nodes1",
+                    cloud_instance_id="0",
+                ),
+                head_node_ip="1.2.3.4",
+            )
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/test_schema.py

@@ -0,0 +1,232 @@
+import os
+import sys
+import time
+
+import pytest
+
+from ray.autoscaler.v2.schema import (
+    AutoscalerInstance,
+    ClusterStatus,
+    IPPRGroupSpec,
+    IPPRStatus,
+)
+from ray.core.generated.autoscaler_pb2 import NodeState, NodeStatus
+from ray.core.generated.instance_manager_pb2 import Instance
+
+
+def test_cluster_status_default_stats():
+    status = ClusterStatus()
+
+    assert status.active_nodes == []
+    assert status.idle_nodes == []
+    assert status.pending_launches == []
+    assert status.failed_launches == []
+    assert status.pending_nodes == []
+    assert status.failed_nodes == []
+    assert status.cluster_resource_usage == []
+    assert status.stats.gcs_request_time_s == 0.0
+    assert status.stats.request_ts_s is None
+
+
+def _make_ippr_status() -> IPPRStatus:
+    return IPPRStatus(
+        cloud_instance_id="ray-worker-1",
+        spec=IPPRGroupSpec(
+            min_cpu=1.0,
+            max_cpu=4.0,
+            min_memory=2,
+            max_memory=8,
+            resize_timeout=10,
+        ),
+        current_cpu=1.0,
+        current_memory=2,
+        desired_cpu=1.0,
+        desired_memory=2,
+    )
+
+
+def test_autoscaler_instance():
+
+    i = AutoscalerInstance()
+    assert not i.validate()[0], "Empty instance should be invalid"
+
+    i = AutoscalerInstance(im_instance=Instance(status=Instance.QUEUED))
+    assert i.validate()[0], "Instance with only im_instance should be valid"
+
+    i = AutoscalerInstance(
+        ray_node=NodeState(status=NodeStatus.RUNNING, instance_id="i-123"),
+        cloud_instance_id="i-123",
+    )
+    assert i.validate()[0], i.validate()[1]
+
+    i = AutoscalerInstance(
+        ray_node=NodeState(status=NodeStatus.RUNNING),
+    )
+    assert not i.validate()[0], "Missing cloud node id."
+
+    i = AutoscalerInstance(
+        im_instance=Instance(status=Instance.QUEUED),
+        ray_node=NodeState(status=NodeStatus.RUNNING),
+    )
+    assert not i.validate()[
+        0
+    ], "cloud node id is required to link the ray node with im state"
+
+    i = AutoscalerInstance(
+        cloud_instance_id="i-123",
+    )
+    assert not i.validate()[
+        0
+    ], "cloud instance id is not possible without im or ray node"
+
+    i = AutoscalerInstance(
+        im_instance=Instance(status=Instance.ALLOCATED, cloud_instance_id="i-123"),
+        cloud_instance_id="i-123",
+    )
+    assert i.validate()[0]
+
+    i = AutoscalerInstance(
+        im_instance=Instance(status=Instance.ALLOCATED, cloud_instance_id="i-123"),
+        cloud_instance_id="i-124",  # mismatch.
+    )
+    assert not i.validate()[0], "cloud instance id should match"
+
+    i = AutoscalerInstance(
+        im_instance=Instance(status=Instance.QUEUED, cloud_instance_id="i-123"),
+        cloud_instance_id="i-123",
+    )
+    assert not i.validate()[0], "cloud instance id is not possible with queued state"
+
+    i = AutoscalerInstance(
+        im_instance=Instance(status=Instance.ALLOCATED, cloud_instance_id="i-123"),
+    )
+    assert not i.validate()[0], "cloud instance id should also be set"
+
+    i = AutoscalerInstance(
+        ray_node=NodeState(status=NodeStatus.RUNNING, instance_id="i-123"),
+        cloud_instance_id="i-123",
+    )
+    assert i.validate()[0]
+
+    i = AutoscalerInstance(
+        ray_node=NodeState(status=NodeStatus.RUNNING, instance_id="i-123"),
+        cloud_instance_id="i-124",  # mismatch.
+    )
+    assert not i.validate()[0]
+
+    i = AutoscalerInstance(
+        im_instance=Instance(status=Instance.RAY_RUNNING, cloud_instance_id="i-123"),
+        ray_node=NodeState(status=NodeStatus.RUNNING, instance_id="i-123"),
+        cloud_instance_id="i-123",
+    )
+    assert i.validate()[0]
+
+    i = AutoscalerInstance(
+        im_instance=Instance(status=Instance.RAY_RUNNING, cloud_instance_id="i-123"),
+        ray_node=NodeState(status=NodeStatus.RUNNING, instance_id="i-123"),
+        cloud_instance_id="i-124",  # mismatch.
+    )
+    assert not i.validate()[0]
+
+
+def test_ippr_status_queue_resize_request():
+    status = _make_ippr_status()
+    status.resizing_at = 123
+    status.k8s_resize_status = "deferred"
+    status.k8s_resize_message = "pending"
+    status.raylet_id = "abc"
+
+    assert status.queue_resize_request(desired_cpu=2.0, desired_memory=4)
+    assert status.desired_cpu == 2.0
+    assert status.desired_memory == 4
+    assert status.resizing_at is None
+    assert status.k8s_resize_status == "new"
+    assert status.k8s_resize_message is None
+
+    assert not status.queue_resize_request(desired_cpu=2.0, desired_memory=4)
+
+
+def test_ippr_status_request_and_progress_helpers():
+    status = _make_ippr_status()
+    assert status.has_resize_request_to_send() is False
+    assert status.is_in_progress() is False
+    status.raylet_id = "abc"
+    assert status.is_k8s_resize_finished()
+    assert not status.has_resize_request_to_send()
+    assert not status.is_in_progress()
+
+    status.queue_resize_request(desired_cpu=2.0)
+    assert status.has_resize_request_to_send()
+    assert status.is_in_progress()
+    assert not status.is_k8s_resize_finished()
+
+    status.resizing_at = int(time.time())
+    status.k8s_resize_status = None
+    assert status.is_in_progress()
+    assert status.is_k8s_resize_finished()
+
+
+def test_ippr_status_need_sync_with_raylet():
+    status = _make_ippr_status()
+    assert status.need_sync_with_raylet() is False
+    status.raylet_id = "abc"
+    status.resizing_at = int(time.time())
+    status.k8s_resize_status = None
+    assert status.need_sync_with_raylet()
+
+    status.desired_cpu = 2.0
+    assert not status.need_sync_with_raylet()
+
+
+def test_ippr_status_limits_and_can_resize_up():
+    status = _make_ippr_status()
+    assert status.can_resize_up()
+    status.raylet_id = "abc"
+    assert status.max_cpu() == 4.0
+    assert status.max_memory() == 8
+    assert status.can_resize_up()
+
+    status.suggested_max_cpu = 1.5
+    status.suggested_max_memory = 3
+    assert status.max_cpu() == 1.5
+    assert status.max_memory() == 3
+
+    status.current_cpu = 1.5
+    status.current_memory = 3
+    assert not status.can_resize_up()
+
+    status.current_cpu = 1.0
+    status.last_failed_at = 1
+    assert not status.can_resize_up()
+
+
+def test_ippr_status_failure_and_timeout_helpers():
+    status = _make_ippr_status()
+    status.raylet_id = "abc"
+    status.desired_cpu = 2.0
+    status.desired_memory = 4
+    status.resizing_at = int(time.time()) - 20
+    assert status.is_timeout()
+
+    status.k8s_resize_status = "error"
+    assert status.is_errored()
+
+    status.record_failure("resize failed", failed_at=123)
+    assert status.last_failed_at == 123
+    assert status.last_failed_reason == "resize failed"
+
+    # K8s resize finished and resources match desired, but raylet sync is still
+    # pending: do not treat as timeout (provider will call GCS to sync).
+    status2 = _make_ippr_status()
+    status2.raylet_id = "abc"
+    status2.resizing_at = int(time.time()) - 20
+    status2.k8s_resize_status = None
+    assert status2.need_sync_with_raylet()
+    assert not status2.is_timeout()
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/test_sdk.py

@@ -0,0 +1,995 @@
+import os
+import sys
+import time
+
+# coding: utf-8
+from dataclasses import dataclass
+from typing import Callable, List, Optional, Tuple
+
+import pytest
+
+import ray
+import ray._private.ray_constants as ray_constants
+from ray._common.test_utils import wait_for_condition
+from ray._private import authentication_test_utils
+from ray.autoscaler.v2.schema import (
+    ClusterStatus,
+    LaunchRequest,
+    NodeInfo,
+    ResourceRequestByCount,
+)
+from ray.autoscaler.v2.sdk import (
+    get_cluster_status,
+    request_cluster_resources,
+)
+from ray.autoscaler.v2.tests.util import (
+    get_available_resources,
+    get_cluster_resource_state,
+    get_total_resources,
+    report_autoscaling_state,
+)
+from ray.core.generated import autoscaler_pb2, autoscaler_pb2_grpc
+from ray.core.generated.autoscaler_pb2 import ClusterResourceState, NodeStatus
+from ray.core.generated.common_pb2 import LabelSelectorOperator
+from ray.util.state.api import list_nodes
+
+
+def _autoscaler_state_service_stub():
+    """Get the grpc stub for the autoscaler state service"""
+    from ray._private.grpc_utils import init_grpc_channel
+
+    gcs_address = ray.get_runtime_context().gcs_address
+    gcs_channel = init_grpc_channel(gcs_address, ray_constants.GLOBAL_GRPC_OPTIONS)
+    return autoscaler_pb2_grpc.AutoscalerStateServiceStub(gcs_channel)
+
+
+def get_node_ids() -> Tuple[str, List[str]]:
+    """Get the node ids of the head node and a worker node"""
+    head_node_id = None
+    nodes = list_nodes()
+    worker_node_ids = []
+    for node in nodes:
+        if node.is_head_node:
+            head_node_id = node.node_id
+        else:
+            worker_node_ids += [node.node_id]
+    return head_node_id, worker_node_ids
+
+
+def assert_cluster_resource_constraints(
+    state: ClusterResourceState, expected_bundles: List[dict], expected_count: List[int]
+):
+    """
+    Assert a GetClusterResourceStateReply has cluster_resource_constraints that
+    matches with the expected resources.
+    """
+    # We only have 1 constraint for now.
+    assert len(state.cluster_resource_constraints) == 1
+
+    resource_requests = state.cluster_resource_constraints[0].resource_requests
+    assert len(resource_requests) == len(expected_bundles) == len(expected_count)
+
+    # Sort all the bundles by bundle's resource names
+    resource_requests = sorted(
+        resource_requests,
+        key=lambda bundle_by_count: "".join(
+            bundle_by_count.request.resources_bundle.keys()
+        ),
+    )
+    expected = zip(expected_bundles, expected_count)
+    expected = sorted(
+        expected, key=lambda bundle_count: "".join(bundle_count[0].keys())
+    )
+
+    for actual_bundle_count, expected_bundle_count in zip(resource_requests, expected):
+        assert (
+            dict(actual_bundle_count.request.resources_bundle)
+            == expected_bundle_count[0]
+        )
+        assert actual_bundle_count.count == expected_bundle_count[1]
+
+
+@dataclass
+class ExpectedNodeState:
+    node_id: str
+    node_status: NodeStatus
+    idle_time_check_cb: Optional[Callable] = None
+    labels: Optional[dict] = None
+
+
+def assert_node_states(
+    state: ClusterResourceState, expected_nodes: List[ExpectedNodeState]
+):
+    """
+    Assert a GetClusterResourceStateReply has node states that
+    matches with the expected nodes.
+    """
+    assert len(state.node_states) == len(expected_nodes)
+
+    # Sort all the nodes by node's node_id
+    node_states = sorted(state.node_states, key=lambda node: node.node_id)
+    expected_nodes = sorted(expected_nodes, key=lambda node: node.node_id)
+
+    for actual_node, expected_node in zip(node_states, expected_nodes):
+        assert actual_node.status == expected_node.node_status
+        if expected_node.idle_time_check_cb:
+            assert expected_node.idle_time_check_cb(actual_node.idle_duration_ms)
+
+        if expected_node.labels:
+            assert sorted(actual_node.dynamic_labels) == sorted(expected_node.labels)
+
+
+@dataclass
+class ExpectedNodeInfo:
+    node_id: Optional[str] = None
+    node_status: Optional[str] = None
+    idle_time_check_cb: Optional[Callable] = None
+    instance_id: Optional[str] = None
+    ray_node_type_name: Optional[str] = None
+    instance_type_name: Optional[str] = None
+    ip_address: Optional[str] = None
+    details: Optional[str] = None
+
+    # Check those resources are included in the actual node info.
+    total_resources: Optional[dict] = None
+    available_resources: Optional[dict] = None
+
+
+def assert_nodes(actual_nodes: List[NodeInfo], expected_nodes: List[ExpectedNodeInfo]):
+
+    assert len(actual_nodes) == len(expected_nodes)
+    # Sort the nodes by id.
+    actual_nodes = sorted(actual_nodes, key=lambda node: node.node_id)
+    expected_nodes = sorted(expected_nodes, key=lambda node: node.node_id)
+
+    for actual_node, expected_node in zip(actual_nodes, expected_nodes):
+        if expected_node.node_id is not None:
+            assert actual_node.node_id == expected_node.node_id
+        if expected_node.node_status is not None:
+            assert actual_node.node_status == expected_node.node_status
+        if expected_node.instance_id is not None:
+            assert actual_node.instance_id == expected_node.instance_id
+        if expected_node.ray_node_type_name is not None:
+            assert actual_node.ray_node_type_name == expected_node.ray_node_type_name
+        if expected_node.instance_type_name is not None:
+            assert actual_node.instance_type_name == expected_node.instance_type_name
+        if expected_node.ip_address is not None:
+            assert actual_node.ip_address == expected_node.ip_address
+        if expected_node.details is not None:
+            assert expected_node.details in actual_node.details
+
+        if expected_node.idle_time_check_cb:
+            assert expected_node.idle_time_check_cb(
+                actual_node.resource_usage.idle_time_ms
+            )
+
+        if expected_node.total_resources:
+            for resource_name, total in expected_node.total_resources.items():
+                assert (
+                    total
+                    == get_total_resources(actual_node.resource_usage.usage)[
+                        resource_name
+                    ]
+                )
+
+        if expected_node.available_resources:
+            for resource_name, available in expected_node.available_resources.items():
+                assert (
+                    available
+                    == get_available_resources(actual_node.resource_usage.usage)[
+                        resource_name
+                    ]
+                )
+
+
+def assert_launches(
+    cluster_status: ClusterStatus,
+    expected_pending_launches: List[LaunchRequest],
+    expected_failed_launches: List[LaunchRequest],
+):
+    def assert_launches(actuals, expects):
+        for actual, expect in zip(actuals, expects):
+            assert actual.instance_type_name == expect.instance_type_name
+            assert actual.ray_node_type_name == expect.ray_node_type_name
+            assert actual.count == expect.count
+            assert actual.state == expect.state
+            assert actual.request_ts_s == expect.request_ts_s
+
+    assert len(cluster_status.pending_launches) == len(expected_pending_launches)
+    assert len(cluster_status.failed_launches) == len(expected_failed_launches)
+
+    actual_pending = sorted(
+        cluster_status.pending_launches, key=lambda launch: launch.ray_node_type_name
+    )
+    expected_pending = sorted(
+        expected_pending_launches, key=lambda launch: launch.ray_node_type_name
+    )
+
+    assert_launches(actual_pending, expected_pending)
+
+    actual_failed = sorted(
+        cluster_status.failed_launches, key=lambda launch: launch.ray_node_type_name
+    )
+    expected_failed = sorted(
+        expected_failed_launches, key=lambda launch: launch.ray_node_type_name
+    )
+
+    assert_launches(actual_failed, expected_failed)
+
+
+@dataclass
+class GangResourceRequest:
+    # Resource bundles.
+    bundles: List[dict]
+    # List of detail information about the request
+    details: List[str]
+
+
+def assert_gang_requests(
+    state: ClusterResourceState, expected: List[GangResourceRequest]
+):
+    """
+    Assert a GetClusterResourceStateReply has gang requests that
+    matches with the expected requests.
+    """
+    assert len(state.pending_gang_resource_requests) == len(expected)
+
+    # Sort all the requests by request's details
+    requests = sorted(
+        state.pending_gang_resource_requests, key=lambda request: request.details
+    )
+    expected = sorted(expected, key=lambda request: "".join(request.details))
+
+    for actual_request, expected_request in zip(requests, expected):
+        # Assert the detail contains the expected details
+        for detail_str in expected_request.details:
+            assert detail_str in actual_request.details
+
+
+def test_request_cluster_resources_basic(shutdown_only):
+    ctx = ray.init(num_cpus=1)
+    stub = _autoscaler_state_service_stub()
+    gcs_address = ctx.address_info["gcs_address"]
+
+    # Request one
+    request_cluster_resources(gcs_address, [{"resources": {"CPU": 1}}])
+
+    def verify():
+        state = get_cluster_resource_state(stub)
+        assert_cluster_resource_constraints(state, [{"CPU": 1}], [1])
+        return True
+
+    wait_for_condition(verify)
+
+    # Request another overrides the previous request
+    request_cluster_resources(
+        gcs_address, [{"resources": {"CPU": 2, "GPU": 1}}, {"resources": {"CPU": 1}}]
+    )
+
+    def verify():
+        state = get_cluster_resource_state(stub)
+        assert_cluster_resource_constraints(
+            state, [{"CPU": 2, "GPU": 1}, {"CPU": 1}], [1, 1]
+        )
+        return True
+
+    # Request multiple is aggregated by shape.
+    request_cluster_resources(gcs_address, [{"resources": {"CPU": 1}}] * 100)
+
+    def verify():
+        state = get_cluster_resource_state(stub)
+        assert_cluster_resource_constraints(state, [{"CPU": 1}], [100])
+        return True
+
+    wait_for_condition(verify)
+
+
+def test_request_cluster_resources_with_label_selectors(shutdown_only):
+    ctx = ray.init(num_cpus=1)
+    stub = _autoscaler_state_service_stub()
+    gcs_address = ctx.address_info["gcs_address"]
+
+    # Define two bundles, each with its own label_selector, to request.
+    bundles = [
+        {"CPU": 1},
+        {"GPU": 1, "CPU": 2},
+    ]
+    bundle_label_selectors = [
+        {"region": "us-west1"},
+        {"accelerator-type": "!in(A100)"},
+    ]
+    to_request = [
+        {"resources": b, "label_selector": s}
+        for b, s in zip(bundles, bundle_label_selectors)
+    ]
+
+    # Send the request for these resource bundles
+    request_cluster_resources(gcs_address, to_request)
+
+    def verify():
+        state = get_cluster_resource_state(stub)
+        # Validate shape and resource request count
+        assert_cluster_resource_constraints(state, bundles, [1, 1])
+
+        # Check that requests carry expected label selectors
+        requests = state.cluster_resource_constraints[0].resource_requests
+
+        # First resource request
+        label_selectors_0 = requests[0].request.label_selectors
+        selector_0 = label_selectors_0[0]
+        constraints_0 = {
+            c.label_key: list(c.label_values) for c in selector_0.label_constraints
+        }
+        assert constraints_0 == {"region": ["us-west1"]}
+        assert (
+            selector_0.label_constraints[0].operator
+            == LabelSelectorOperator.LABEL_OPERATOR_IN
+        )
+
+        # Second resource request
+        label_selectors_1 = requests[1].request.label_selectors
+        selector_1 = label_selectors_1[0]
+        constraints_1 = {
+            c.label_key: list(c.label_values) for c in selector_1.label_constraints
+        }
+        assert constraints_1 == {"accelerator-type": ["A100"]}
+        assert (
+            selector_1.label_constraints[0].operator
+            == LabelSelectorOperator.LABEL_OPERATOR_NOT_IN
+        )
+
+        return True
+
+    wait_for_condition(verify)
+
+
+def test_node_info_basic(shutdown_only, monkeypatch):
+    with monkeypatch.context() as m:
+        m.setenv("RAY_CLOUD_INSTANCE_ID", "instance-id")
+        m.setenv("RAY_NODE_TYPE_NAME", "node-type-name")
+        m.setenv("RAY_CLOUD_INSTANCE_TYPE_NAME", "instance-type-name")
+
+        ctx = ray.init(num_cpus=1)
+        ip = ctx.address_info["node_ip_address"]
+
+        stub = _autoscaler_state_service_stub()
+
+        def verify():
+            state = get_cluster_resource_state(stub)
+
+            assert len(state.node_states) == 1
+            node = state.node_states[0]
+
+            assert node.instance_id == "instance-id"
+            assert node.ray_node_type_name == "node-type-name"
+            assert node.node_ip_address == ip
+            assert node.instance_type_name == "instance-type-name"
+
+            assert (
+                state.cluster_session_name
+                == ray._private.worker.global_worker.node.session_name
+            )
+            return True
+
+        wait_for_condition(verify)
+
+
+def test_pg_pending_gang_requests_basic(shutdown_only):
+    ray.init(num_cpus=1)
+
+    # Create a pg that's pending.
+    pg = ray.util.placement_group([{"CPU": 1}] * 3, strategy="STRICT_SPREAD")
+    try:
+        ray.get(pg.ready(), timeout=2)
+    except TimeoutError:
+        pass
+
+    pg_id = pg.id.hex()
+
+    stub = _autoscaler_state_service_stub()
+
+    def verify():
+        state = get_cluster_resource_state(stub)
+        assert_gang_requests(
+            state,
+            [
+                GangResourceRequest(
+                    [{"CPU": 1}] * 3, details=[pg_id, "STRICT_SPREAD", "PENDING"]
+                )
+            ],
+        )
+        return True
+
+    wait_for_condition(verify)
+
+
+def test_pg_usage_labels(shutdown_only):
+    ray.init(num_cpus=1)
+
+    # Create a pg
+    pg = ray.util.placement_group([{"CPU": 1}])
+    ray.get(pg.ready())
+
+    # Check the labels
+    stub = _autoscaler_state_service_stub()
+    head_node_id, _ = get_node_ids()
+
+    pg_id = pg.id.hex()
+
+    def verify():
+        state = get_cluster_resource_state(stub)
+        assert_node_states(
+            state,
+            [
+                ExpectedNodeState(
+                    head_node_id,
+                    NodeStatus.RUNNING,
+                    labels={f"_PG_{pg_id}": ""},
+                ),
+            ],
+        )
+        return True
+
+    wait_for_condition(verify)
+
+
+def test_node_state_lifecycle_basic(ray_start_cluster):
+    start_s = time.perf_counter()
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=0)
+    ray.init(address=cluster.address)
+
+    node = cluster.add_node(num_cpus=1)
+
+    stub = _autoscaler_state_service_stub()
+
+    # We don't have node id from `add_node` unfortunately.
+    def nodes_up():
+        nodes = list_nodes()
+        assert len(nodes) == 2
+        return True
+
+    wait_for_condition(nodes_up)
+
+    head_node_id, worker_node_ids = get_node_ids()
+    node_id = worker_node_ids[0]
+
+    def verify_cluster_idle():
+        state = get_cluster_resource_state(stub)
+        assert_node_states(
+            state,
+            [
+                ExpectedNodeState(
+                    node_id, NodeStatus.IDLE, lambda idle_ms: idle_ms > 0
+                ),
+                ExpectedNodeState(
+                    head_node_id, NodeStatus.IDLE, lambda idle_ms: idle_ms > 0
+                ),
+            ],
+        )
+
+        return True
+
+    wait_for_condition(verify_cluster_idle)
+
+    # Schedule a task running
+    @ray.remote(num_cpus=0.1)
+    def f():
+        while True:
+            pass
+
+    t = f.remote()
+
+    def verify_cluster_busy():
+        state = get_cluster_resource_state(stub)
+        assert_node_states(
+            state,
+            [
+                ExpectedNodeState(
+                    node_id, NodeStatus.RUNNING, lambda idle_ms: idle_ms == 0
+                ),
+                ExpectedNodeState(
+                    head_node_id, NodeStatus.IDLE, lambda idle_ms: idle_ms > 0
+                ),
+            ],
+        )
+        return True
+
+    wait_for_condition(verify_cluster_busy)
+
+    # Kill the task
+    ray.cancel(t, force=True)
+
+    wait_for_condition(verify_cluster_idle)
+
+    # Kill the node.
+    cluster.remove_node(node)
+
+    # Sleep for a bit so head node should be idle longer than this.
+    time.sleep(3)
+
+    def verify_cluster_no_node():
+        state = get_cluster_resource_state(stub)
+        now_s = time.perf_counter()
+        test_dur_ms = (now_s - start_s) * 1000
+        assert_node_states(
+            state,
+            [
+                ExpectedNodeState(node_id, NodeStatus.DEAD),
+                ExpectedNodeState(
+                    head_node_id,
+                    NodeStatus.IDLE,
+                    lambda idle_ms: idle_ms > 3 * 1000 and idle_ms < test_dur_ms,
+                ),
+            ],
+        )
+        return True
+
+    wait_for_condition(verify_cluster_no_node)
+
+
+# We test that a node with only workers blocked on get
+# is considered idle.
+def test_idle_node_blocked(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=1)
+    ray.init(address=cluster.address)
+
+    stub = _autoscaler_state_service_stub()
+
+    # We don't have node id from `add_node` unfortunately.
+    def nodes_up():
+        nodes = list_nodes()
+        assert len(nodes) == 1
+        return True
+
+    wait_for_condition(nodes_up)
+
+    head_node_id = get_node_ids()
+
+    def verify_cluster_idle():
+        state = get_cluster_resource_state(stub)
+        assert_node_states(
+            state,
+            [
+                ExpectedNodeState(
+                    head_node_id, NodeStatus.IDLE, lambda idle_ms: idle_ms > 0
+                ),
+            ],
+        )
+        return True
+
+    wait_for_condition(verify_cluster_idle)
+
+    # Unschedulable
+    @ray.remote(num_cpus=10000)
+    def f():
+        pass
+
+    # Schedule a task running
+    @ray.remote(num_cpus=1)
+    def g():
+        ray.get(f.remote())
+
+    t = g.remote()
+
+    def verify_cluster_busy():
+        state = get_cluster_resource_state(stub)
+        assert_node_states(
+            state,
+            [
+                ExpectedNodeState(
+                    head_node_id, NodeStatus.RUNNING, lambda idle_ms: idle_ms == 0
+                ),
+            ],
+        )
+        return True
+
+    wait_for_condition(verify_cluster_busy)
+    for _ in range(10):
+        time.sleep(0.5)
+        verify_cluster_busy()
+
+    # Kill the task
+    ray.cancel(t, force=True)
+    wait_for_condition(verify_cluster_idle)
+
+
+def test_idle_node_no_resource(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=1)
+    ray.init(address=cluster.address)
+
+    stub = _autoscaler_state_service_stub()
+
+    # We don't have node id from `add_node` unfortunately.
+    def nodes_up():
+        nodes = list_nodes()
+        assert len(nodes) == 1
+        return True
+
+    wait_for_condition(nodes_up)
+
+    head_node_id = get_node_ids()
+
+    def verify_cluster_idle():
+        state = get_cluster_resource_state(stub)
+        assert_node_states(
+            state,
+            [
+                ExpectedNodeState(
+                    head_node_id, NodeStatus.IDLE, lambda idle_ms: idle_ms > 0
+                ),
+            ],
+        )
+        return True
+
+    wait_for_condition(verify_cluster_idle)
+
+    # Schedule a task running
+    @ray.remote(num_cpus=0)
+    def f():
+        while True:
+            pass
+
+    t = f.remote()
+
+    def verify_cluster_busy():
+        state = get_cluster_resource_state(stub)
+        assert_node_states(
+            state,
+            [
+                ExpectedNodeState(
+                    head_node_id, NodeStatus.RUNNING, lambda idle_ms: idle_ms == 0
+                ),
+            ],
+        )
+        return True
+
+    wait_for_condition(verify_cluster_busy)
+
+    # Kill the task
+    ray.cancel(t, force=True)
+    wait_for_condition(verify_cluster_idle)
+
+
+def test_get_cluster_status_resources(ray_start_cluster):
+    cluster = ray_start_cluster
+    # Head node
+    cluster.add_node(num_cpus=1, _system_config={"enable_autoscaler_v2": True})
+    ray.init(address=cluster.address)
+
+    # Worker node
+    cluster.add_node(num_cpus=2)
+
+    @ray.remote(num_cpus=1)
+    class Actor:
+        def loop(self):
+            while True:
+                pass
+
+    # Schedule tasks to use all resources.
+    @ray.remote(num_cpus=1)
+    def loop():
+        while True:
+            pass
+
+    [loop.remote() for _ in range(2)]
+    actor = Actor.remote()
+    actor.loop.remote()
+
+    def verify_cpu_resources_all_used():
+        cluster_status = get_cluster_status(cluster.address)
+        total_cluster_resources = get_total_resources(
+            cluster_status.cluster_resource_usage
+        )
+        assert total_cluster_resources["CPU"] == 3.0
+
+        available_cluster_resources = get_available_resources(
+            cluster_status.cluster_resource_usage
+        )
+        assert available_cluster_resources["CPU"] == 0.0
+        return True
+
+    wait_for_condition(verify_cpu_resources_all_used)
+
+    # Schedule more tasks should show up as task demands
+    [loop.remote() for _ in range(2)]
+
+    def verify_task_demands():
+        resource_demands = get_cluster_status(cluster.address).resource_demands
+        assert len(resource_demands.ray_task_actor_demand) == 1
+        assert resource_demands.ray_task_actor_demand[0].bundles_by_count == [
+            ResourceRequestByCount(
+                bundle={"CPU": 1.0},
+                count=2,
+            )
+        ]
+        return True
+
+    wait_for_condition(verify_task_demands)
+
+    # Request resources through SDK
+    request_cluster_resources(
+        gcs_address=cluster.address, to_request=[{"resources": {"GPU": 1, "CPU": 2}}]
+    )
+
+    def verify_cluster_constraint_demand():
+        resource_demands = get_cluster_status(cluster.address).resource_demands
+        assert len(resource_demands.cluster_constraint_demand) == 1
+        assert resource_demands.cluster_constraint_demand[0].bundles_by_count == [
+            ResourceRequestByCount(
+                bundle={"GPU": 1.0, "CPU": 2.0},
+                count=1,
+            )
+        ]
+        return True
+
+    wait_for_condition(verify_cluster_constraint_demand)
+
+    # Try to schedule some PGs
+    pg1 = ray.util.placement_group([{"CPU": 1}] * 3)
+
+    def verify_pg_demands():
+        resource_demands = get_cluster_status(cluster.address).resource_demands
+        assert len(resource_demands.placement_group_demand) == 1
+        assert resource_demands.placement_group_demand[0].bundles_by_count == [
+            ResourceRequestByCount(
+                bundle={"CPU": 1.0},
+                count=3,
+            )
+        ]
+        assert resource_demands.placement_group_demand[0].pg_id == pg1.id.hex()
+        assert resource_demands.placement_group_demand[0].strategy == "PACK"
+        assert resource_demands.placement_group_demand[0].state == "PENDING"
+        return True
+
+    wait_for_condition(verify_pg_demands)
+
+
+def test_get_cluster_status(ray_start_cluster):
+    # This test is to make sure the grpc stub is working.
+    # TODO(rickyx): Add e2e tests for the autoscaler state service in a separate PR
+    # to validate the data content.
+    cluster = ray_start_cluster
+    # Head node
+    cluster.add_node(num_cpus=1, _system_config={"enable_autoscaler_v2": True})
+    ray.init(address=cluster.address)
+    # Worker node
+    cluster.add_node(num_cpus=2)
+
+    head_node_id, worker_node_ids = get_node_ids()
+
+    def verify_nodes():
+        cluster_status = get_cluster_status(cluster.address)
+        assert_nodes(
+            cluster_status.idle_nodes,
+            [
+                ExpectedNodeInfo(
+                    worker_node_ids[0],
+                    "IDLE",
+                    lambda idle_ms: idle_ms > 0,
+                    total_resources={"CPU": 2.0},
+                    available_resources={"CPU": 2.0},
+                ),
+                ExpectedNodeInfo(
+                    head_node_id,
+                    "IDLE",
+                    lambda idle_ms: idle_ms > 0,
+                    total_resources={"CPU": 1.0},
+                    available_resources={"CPU": 1.0},
+                ),
+            ],
+        )
+        return True
+
+    wait_for_condition(verify_nodes)
+
+    # Schedule a task running
+    @ray.remote(num_cpus=2)
+    def f():
+        while True:
+            pass
+
+    f.remote()
+
+    def verify_nodes_busy():
+        cluster_status = get_cluster_status(cluster.address)
+        assert_nodes(
+            cluster_status.idle_nodes,
+            [
+                ExpectedNodeInfo(head_node_id, "IDLE", lambda idle_ms: idle_ms > 0),
+            ],
+        )
+        assert_nodes(
+            cluster_status.active_nodes,
+            [
+                ExpectedNodeInfo(
+                    worker_node_ids[0],
+                    "RUNNING",
+                    lambda idle_ms: idle_ms == 0,
+                    total_resources={"CPU": 2.0},
+                    available_resources={"CPU": 0.0},
+                ),
+            ],
+        )
+        return True
+
+    wait_for_condition(verify_nodes_busy)
+
+    stub = _autoscaler_state_service_stub()
+    state = autoscaler_pb2.AutoscalingState(
+        last_seen_cluster_resource_state_version=0,
+        # since the autoscaler will also update the autoscaler_state_version periodically,
+        # we need to use a large number here, such as 10, to override it to avoid flaky test.
+        autoscaler_state_version=10,
+        pending_instance_requests=[
+            autoscaler_pb2.PendingInstanceRequest(
+                instance_type_name="m5.large",
+                ray_node_type_name="worker",
+                count=2,
+                request_ts=1000,
+            )
+        ],
+        failed_instance_requests=[
+            autoscaler_pb2.FailedInstanceRequest(
+                instance_type_name="m5.large",
+                ray_node_type_name="worker",
+                count=2,
+                start_ts=1000,
+                failed_ts=2000,
+                reason="insufficient quota",
+            )
+        ],
+        pending_instances=[
+            autoscaler_pb2.PendingInstance(
+                instance_id="instance-id",
+                instance_type_name="m5.large",
+                ray_node_type_name="worker",
+                ip_address="10.10.10.10",
+                details="launching",
+            )
+        ],
+    )
+    report_autoscaling_state(stub, autoscaling_state=state)
+
+    def verify_autoscaler_state():
+        # TODO(rickyx): Add infeasible asserts.
+
+        cluster_status = get_cluster_status(cluster.address)
+        assert len(cluster_status.pending_launches) == 1
+        assert_launches(
+            cluster_status,
+            expected_pending_launches=[
+                LaunchRequest(
+                    instance_type_name="m5.large",
+                    ray_node_type_name="worker",
+                    count=2,
+                    state=LaunchRequest.Status.PENDING,
+                    request_ts_s=1000,
+                )
+            ],
+            expected_failed_launches=[
+                LaunchRequest(
+                    instance_type_name="m5.large",
+                    ray_node_type_name="worker",
+                    count=2,
+                    state=LaunchRequest.Status.FAILED,
+                    request_ts_s=1000,
+                    failed_ts_s=2000,
+                    details="insufficient quota",
+                )
+            ],
+        )
+
+        assert_nodes(
+            cluster_status.pending_nodes,
+            [
+                ExpectedNodeInfo(
+                    instance_id="instance-id",
+                    ray_node_type_name="worker",
+                    details="launching",
+                    ip_address="10.10.10.10",
+                )
+            ],
+        )
+        return True
+
+    wait_for_condition(verify_autoscaler_state)
+
+
+@pytest.mark.parametrize(
+    "env_val,enabled",
+    [
+        ("1", True),
+        ("0", False),
+        ("", False),
+    ],
+)
+def test_is_autoscaler_v2_enabled(shutdown_only, monkeypatch, env_val, enabled):
+    def reset_autoscaler_v2_enabled_cache():
+        import ray.autoscaler.v2.utils as u
+
+        u.cached_is_autoscaler_v2 = None
+
+    reset_autoscaler_v2_enabled_cache()
+    with monkeypatch.context() as m:
+        m.setenv("RAY_enable_autoscaler_v2", env_val)
+        ray.init()
+
+        def verify():
+            assert ray.autoscaler.v2.utils.is_autoscaler_v2() == enabled
+            return True
+
+        wait_for_condition(verify)
+
+
+@pytest.mark.parametrize(
+    "token_state,setup_token,should_fail",
+    [
+        ("valid", lambda: None, False),
+        ("invalid", lambda: _setup_invalid_token(), True),
+    ],
+)
+def test_autoscaler_api_with_token_auth(
+    setup_cluster_with_token_auth,
+    cleanup_auth_token_env,
+    token_state,
+    setup_token,
+    should_fail,
+):
+    """Parametrized test for autoscaler API with different token states.
+
+    Tests request_cluster_resources with valid, invalid, and missing tokens.
+    """
+    # Setup token state (this changes the client-side token)
+    setup_token()
+
+    if should_fail:
+        # API call should fail with invalid token
+        with pytest.raises(Exception) as exc_info:
+            request_cluster_resources(
+                ray.get_runtime_context().gcs_address,
+                [{"resources": {"CPU": 1}, "label_selector": {}}],
+            )
+
+        # Verify it's an authentication error
+        error_str = str(exc_info.value).lower()
+        assert (
+            "unauthenticated" in error_str or "invalidauthtoken" in error_str
+        ), f"request_cluster_resources with {token_state} token should return auth error, got: {exc_info.value}"
+    else:
+        # API call should succeed with valid token
+        request_cluster_resources(
+            ray.get_runtime_context().gcs_address,
+            [{"resources": {"CPU": 1}, "label_selector": {}}],
+        )
+
+        # Verify the request was successful using the autoscaler state service stub
+        stub = _autoscaler_state_service_stub()
+        state = get_cluster_resource_state(stub)
+        assert (
+            len(state.cluster_resource_constraints) > 0
+        ), f"request_cluster_resources with {token_state} token should succeed"
+
+
+def _setup_invalid_token():
+    """Helper to set up an invalid authentication token."""
+
+    invalid_token = "invalid_token_value"
+    authentication_test_utils.set_env_auth_token(invalid_token)
+    authentication_test_utils.reset_auth_token_state()
+
+
+def _clear_token():
+    """Helper to clear authentication token sources."""
+
+    authentication_test_utils.clear_auth_token_sources()
+    authentication_test_utils.reset_auth_token_state()
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/test_storage.py

@@ -0,0 +1,87 @@
+# coding: utf-8
+import os
+import sys
+
+import pytest  # noqa
+
+from ray.autoscaler.v2.instance_manager.storage import (
+    InMemoryStorage,
+    StoreStatus,
+    VersionedValue,
+)
+
+
+@pytest.mark.parametrize("storage", [InMemoryStorage()])
+def test_storage(storage):
+    assert storage.get_version() == 0
+    assert storage.get_all(table="test_table") == ({}, 0)
+    assert storage.get(table="test_table", keys=[]) == ({}, 0)
+    assert storage.get(table="test_table", keys=["key1"]) == ({}, 0)
+
+    assert storage.batch_update(
+        table="test_table", mutation={"key1": "value1"}
+    ) == StoreStatus(
+        True,
+        1,
+    )
+
+    assert storage.get_version() == 1
+
+    assert storage.get_all(table="test_table") == (
+        {"key1": VersionedValue("value1", 1)},
+        1,
+    )
+    assert storage.get(table="test_table", keys=[]) == (
+        {"key1": VersionedValue("value1", 1)},
+        1,
+    )
+
+    assert storage.batch_update(
+        table="test_table", mutation={"key1": "value2"}, expected_version=0
+    ) == StoreStatus(False, 1)
+
+    assert storage.batch_update(
+        table="test_table", mutation={"key1": "value2"}, expected_version=1
+    ) == StoreStatus(True, 2)
+
+    assert storage.get_all(table="test_table") == (
+        {"key1": VersionedValue("value2", 2)},
+        2,
+    )
+
+    assert storage.batch_update(
+        table="test_table",
+        mutation={"key2": "value3", "key3": "value4"},
+        deletion=["key1"],
+        expected_version=2,
+    ) == StoreStatus(True, 3)
+
+    assert storage.get_all(table="test_table") == (
+        {"key2": VersionedValue("value3", 3), "key3": VersionedValue("value4", 3)},
+        3,
+    )
+
+    assert storage.get(table="test_table", keys=["key2", "key1"]) == (
+        {"key2": VersionedValue("value3", 3)},
+        3,
+    )
+
+    assert storage.update(
+        table="test_table", key="key2", value="value5"
+    ) == StoreStatus(True, 4)
+    assert storage.update(
+        table="test_table", key="key2", value="value5", insert_only=True
+    ) == StoreStatus(False, 4)
+    assert storage.update(
+        table="test_table", key="key2", value="value5", expected_entry_version=3
+    ) == StoreStatus(False, 4)
+    assert storage.update(
+        table="test_table", key="key2", value="value6", expected_entry_version=4
+    ) == StoreStatus(True, 5)
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/test_subscribers.py

@@ -0,0 +1,392 @@
+# coding: utf-8
+import os
+import sys
+from queue import Queue
+from unittest import mock
+
+import pytest
+
+from ray._common.test_utils import wait_for_condition
+from ray._common.utils import binary_to_hex, hex_to_binary
+from ray.autoscaler._private.prom_metrics import AutoscalerPrometheusMetrics
+from ray.autoscaler.v2.instance_manager.cloud_providers.read_only.cloud_provider import (
+    ReadOnlyProvider,
+)
+from ray.autoscaler.v2.instance_manager.subscribers.cloud_instance_updater import (
+    CloudInstanceUpdater,
+)
+from ray.autoscaler.v2.instance_manager.subscribers.ray_stopper import (  # noqa
+    RayStopper,
+)
+from ray.autoscaler.v2.metrics_reporter import AutoscalerMetricsReporter
+from ray.core.generated.autoscaler_pb2 import DrainNodeReason
+from ray.core.generated.instance_manager_pb2 import (
+    Instance,
+    InstanceUpdateEvent,
+    TerminationRequest,
+)
+
+
+def _get_stopped_nodes_total(metrics_reporter: AutoscalerMetricsReporter) -> float:
+    total_samples = [
+        sample.value
+        for metric in metrics_reporter._prom_metrics.stopped_nodes.collect()
+        for sample in metric.samples
+        if sample.name == "autoscaler_stopped_nodes_total"
+    ]
+    assert len(total_samples) == 1
+    return total_samples[0]
+
+
+class TestRayStopper:
+    def test_no_op(self):
+        mock_gcs_client = mock.MagicMock()
+        ray_stopper = RayStopper(gcs_client=mock_gcs_client, error_queue=Queue())
+
+        ray_stopper.notify(
+            [
+                InstanceUpdateEvent(
+                    instance_id="test_id",
+                    new_instance_status=Instance.REQUESTED,
+                )
+            ]
+        )
+        assert mock_gcs_client.drain_node.call_count == 0
+
+    @pytest.mark.parametrize(
+        "drain_accepted",
+        [True, False, None],
+        ids=["drain_accepted", "drain_rejected", "drain_error"],
+    )
+    def test_idle_termination(self, drain_accepted):
+        mock_gcs_client = mock.MagicMock()
+        if drain_accepted is None:
+            mock_gcs_client.drain_node.side_effect = Exception("error")
+        else:
+            mock_gcs_client.drain_node.return_value = (
+                drain_accepted,
+                f"accepted={str(drain_accepted)}",
+            )
+        error_queue = Queue()
+        ray_stopper = RayStopper(gcs_client=mock_gcs_client, error_queue=error_queue)
+
+        ray_stopper.notify(
+            [
+                InstanceUpdateEvent(
+                    instance_id="test_id",
+                    new_instance_status=Instance.RAY_STOP_REQUESTED,
+                    termination_request=TerminationRequest(
+                        cause=TerminationRequest.Cause.IDLE,
+                        idle_duration_ms=1000,
+                        ray_node_id="0000",
+                    ),
+                )
+            ]
+        )
+
+        def verify():
+            mock_gcs_client.drain_node.assert_has_calls(
+                [
+                    mock.call(
+                        node_id="0000",
+                        reason=DrainNodeReason.DRAIN_NODE_REASON_IDLE_TERMINATION,
+                        reason_message=(
+                            "Termination of node that's idle for 1.0 seconds."
+                        ),
+                        deadline_timestamp_ms=0,
+                    )
+                ]
+            )
+
+            if drain_accepted:
+                assert error_queue.empty()
+            else:
+                error = error_queue.get_nowait()
+                assert error.im_instance_id == "test_id"
+
+            return True
+
+        wait_for_condition(verify)
+
+    @pytest.mark.parametrize(
+        "stop_accepted",
+        [True, False],
+        ids=["stop_accepted", "stop_rejected"],
+    )
+    def test_preemption(self, stop_accepted):
+        mock_gcs_client = mock.MagicMock()
+        mock_gcs_client.drain_nodes.return_value = [0] if stop_accepted else []
+        error_queue = Queue()
+        ray_stopper = RayStopper(gcs_client=mock_gcs_client, error_queue=error_queue)
+
+        ray_stopper.notify(
+            [
+                InstanceUpdateEvent(
+                    instance_id="i-1",
+                    new_instance_status=Instance.RAY_STOP_REQUESTED,
+                    termination_request=TerminationRequest(
+                        cause=TerminationRequest.Cause.MAX_NUM_NODE_PER_TYPE,
+                        max_num_nodes_per_type=10,
+                        ray_node_id=binary_to_hex(hex_to_binary(b"1111")),
+                    ),
+                ),
+                InstanceUpdateEvent(
+                    instance_id="i-2",
+                    new_instance_status=Instance.RAY_STOP_REQUESTED,
+                    termination_request=TerminationRequest(
+                        cause=TerminationRequest.Cause.MAX_NUM_NODES,
+                        max_num_nodes=100,
+                        ray_node_id=binary_to_hex(hex_to_binary(b"2222")),
+                    ),
+                ),
+            ]
+        )
+
+        def verify():
+            mock_gcs_client.drain_nodes.assert_has_calls(
+                [
+                    mock.call(
+                        node_ids=[hex_to_binary(b"1111")],
+                    ),
+                    mock.call(
+                        node_ids=[hex_to_binary(b"2222")],
+                    ),
+                ]
+            )
+
+            if stop_accepted:
+                assert error_queue.empty()
+            else:
+                error_in_ids = set()
+                while not error_queue.empty():
+                    error = error_queue.get_nowait()
+                    error_in_ids.add(error.im_instance_id)
+
+                assert error_in_ids == {"i-1", "i-2"}
+
+            return True
+
+        wait_for_condition(verify)
+
+
+class TestCloudInstanceUpdater:
+    def test_launch_no_op(self):
+        mock_provider = mock.MagicMock()
+        launcher = CloudInstanceUpdater(mock_provider)
+        launcher.notify(
+            [
+                InstanceUpdateEvent(
+                    new_instance_status=Instance.RAY_RUNNING,
+                    launch_request_id="1",
+                    instance_type="type-1",
+                ),
+            ]
+        )
+        mock_provider.launch.assert_not_called()
+
+    def test_launch_new_instances(self):
+        mock_provider = mock.MagicMock()
+        launcher = CloudInstanceUpdater(mock_provider)
+        launcher.notify(
+            [
+                InstanceUpdateEvent(
+                    new_instance_status=Instance.REQUESTED,
+                    launch_request_id="1",
+                    instance_type="type-1",
+                ),
+                InstanceUpdateEvent(
+                    new_instance_status=Instance.REQUESTED,
+                    launch_request_id="1",
+                    instance_type="type-1",
+                ),
+                InstanceUpdateEvent(
+                    new_instance_status=Instance.REQUESTED,
+                    launch_request_id="2",
+                    instance_type="type-1",
+                ),
+                InstanceUpdateEvent(
+                    new_instance_status=Instance.REQUESTED,
+                    launch_request_id="2",
+                    instance_type="type-2",
+                ),
+            ]
+        )
+
+        def verify():
+            mock_provider.launch.assert_has_calls(
+                [
+                    mock.call(shape={"type-1": 2}, request_id="1"),
+                    mock.call(shape={"type-1": 1, "type-2": 1}, request_id="2"),
+                ]
+            )
+            return True
+
+        wait_for_condition(verify)
+
+    def test_multi_notify(self):
+        mock_provider = mock.MagicMock()
+        launcher = CloudInstanceUpdater(mock_provider)
+
+        launcher.notify(
+            [
+                InstanceUpdateEvent(
+                    new_instance_status=Instance.REQUESTED,
+                    launch_request_id="1",
+                    instance_type="type-1",
+                ),
+            ]
+        )
+
+        launcher.notify(
+            [
+                InstanceUpdateEvent(
+                    new_instance_status=Instance.REQUESTED,
+                    launch_request_id="2",
+                    instance_type="type-1",
+                ),
+            ]
+        )
+
+        def verify():
+            assert mock_provider.launch.call_count == 2
+            mock_provider.launch.assert_has_calls(
+                [
+                    mock.call(shape={"type-1": 1}, request_id="1"),
+                    mock.call(shape={"type-1": 1}, request_id="2"),
+                ]
+            )
+            return True
+
+        wait_for_condition(verify)
+
+    def test_terminate_no_op(self):
+        mock_provider = mock.MagicMock()
+        launcher = CloudInstanceUpdater(mock_provider)
+        launcher.notify(
+            [
+                InstanceUpdateEvent(
+                    new_instance_status=Instance.RAY_RUNNING,
+                    instance_id="1",
+                    cloud_instance_id="c1",
+                ),
+            ]
+        )
+
+        def verify():
+            mock_provider.terminate.assert_not_called()
+            return True
+
+        wait_for_condition(verify)
+
+    def test_terminate_instances(self):
+        mock_provider = mock.MagicMock()
+        metrics_reporter = AutoscalerMetricsReporter(
+            AutoscalerPrometheusMetrics(session_name="test")
+        )
+        launcher = CloudInstanceUpdater(mock_provider, metrics_reporter)
+        launcher.notify(
+            [
+                InstanceUpdateEvent(
+                    new_instance_status=Instance.TERMINATING,
+                    instance_id="1",
+                    cloud_instance_id="c1",
+                ),
+                InstanceUpdateEvent(
+                    new_instance_status=Instance.TERMINATING,
+                    instance_id="2",
+                    cloud_instance_id="c2",
+                ),
+                InstanceUpdateEvent(
+                    new_instance_status=Instance.TERMINATING,
+                    instance_id="3",
+                    cloud_instance_id="c3",
+                ),
+            ]
+        )
+
+        def verify():
+            mock_provider.terminate.assert_called_once_with(
+                ids=["c1", "c2", "c3"], request_id=mock.ANY
+            )
+            assert _get_stopped_nodes_total(metrics_reporter) == 0
+            return True
+
+        wait_for_condition(verify)
+
+    def test_count_stopped_instances_on_terminated(self):
+        mock_provider = mock.MagicMock()
+        metrics_reporter = AutoscalerMetricsReporter(
+            AutoscalerPrometheusMetrics(session_name="test")
+        )
+        launcher = CloudInstanceUpdater(mock_provider, metrics_reporter)
+        launcher.notify(
+            [
+                InstanceUpdateEvent(
+                    new_instance_status=Instance.TERMINATED,
+                    instance_id="1",
+                    cloud_instance_id="c1",
+                ),
+                InstanceUpdateEvent(
+                    new_instance_status=Instance.TERMINATED,
+                    instance_id="2",
+                    cloud_instance_id="c2",
+                ),
+                InstanceUpdateEvent(
+                    new_instance_status=Instance.TERMINATED,
+                    instance_id="3",
+                ),
+            ]
+        )
+
+        def verify():
+            mock_provider.terminate.assert_not_called()
+            assert _get_stopped_nodes_total(metrics_reporter) == 2
+            return True
+
+        wait_for_condition(verify)
+
+
+class TestReadOnlyProvider:
+    def test_terminate_raises_not_implemented_with_correct_interface(self):
+        """ReadOnlyProvider.terminate() must accept ids and request_id kwargs.
+
+        Regression test for:
+          TypeError: ReadOnlyProvider.terminate() got an unexpected keyword
+          argument 'ids'
+
+        CloudInstanceUpdater calls provider.terminate(ids=..., request_id=...)
+        which matches the ICloudInstanceProvider interface. ReadOnlyProvider
+        must accept the same signature even though it raises NotImplementedError.
+        """
+        provider = object.__new__(ReadOnlyProvider)  # skip __init__ (needs GCS)
+
+        with pytest.raises(NotImplementedError):
+            provider.terminate(ids=["node-1", "node-2"], request_id="req-1")
+
+    def test_terminate_via_cloud_instance_updater_raises_not_implemented(self):
+        """Verify the full call path: CloudInstanceUpdater -> ReadOnlyProvider.
+
+        When a TERMINATING event arrives, CloudInstanceUpdater calls
+        provider.terminate(ids=..., request_id=...). With ReadOnlyProvider this
+        should surface as NotImplementedError, not TypeError.
+        """
+        provider = object.__new__(ReadOnlyProvider)
+        updater = CloudInstanceUpdater(cloud_provider=provider)
+
+        with pytest.raises(NotImplementedError):
+            updater.notify(
+                [
+                    InstanceUpdateEvent(
+                        new_instance_status=Instance.TERMINATING,
+                        instance_id="i-1",
+                        cloud_instance_id="cloud-node-1",
+                    ),
+                ]
+            )
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/test_threaded_ray_installer.py

@@ -0,0 +1,119 @@
+# coding: utf-8
+import os
+import sys
+import unittest
+from queue import Queue
+from unittest.mock import patch
+
+import pytest  # noqa
+
+from ray._private.test_utils import load_test_config
+from ray.autoscaler.tags import TAG_RAY_NODE_KIND
+from ray.autoscaler.v2.instance_manager.config import AutoscalingConfig
+from ray.autoscaler.v2.instance_manager.instance_storage import InstanceStorage
+from ray.autoscaler.v2.instance_manager.ray_installer import RayInstaller
+from ray.autoscaler.v2.instance_manager.storage import InMemoryStorage
+from ray.autoscaler.v2.instance_manager.subscribers.threaded_ray_installer import (
+    ThreadedRayInstaller,
+)
+from ray.core.generated.instance_manager_pb2 import Instance, NodeKind
+from ray.tests.autoscaler_test_utils import MockProcessRunner, MockProvider
+
+
+class ThreadedRayInstallerTest(unittest.TestCase):
+    def setUp(self):
+        self.base_provider = MockProvider()
+        self.config = AutoscalingConfig(load_test_config("test_ray_complex.yaml"))
+        self.runner = MockProcessRunner()
+        self.ray_installer = RayInstaller(self.base_provider, self.config, self.runner)
+        self.instance_storage = InstanceStorage(
+            cluster_id="test_cluster_id",
+            storage=InMemoryStorage(),
+        )
+        self.error_queue = Queue()
+        self.threaded_ray_installer = ThreadedRayInstaller(
+            head_node_ip="127.0.0.1",
+            instance_storage=self.instance_storage,
+            ray_installer=self.ray_installer,
+            error_queue=self.error_queue,
+        )
+
+    def test_install_ray_on_new_node_version_mismatch(self):
+        self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
+        instance = Instance(
+            instance_id="0",
+            instance_type="worker_nodes1",
+            cloud_instance_id="0",
+            status=Instance.RAY_INSTALLING,
+            node_kind=NodeKind.WORKER,
+        )
+        success, verison = self.instance_storage.upsert_instance(instance)
+        assert success
+        self.runner.respond_to_call("json .Config.Env", ["[]" for i in range(1)])
+
+        self.threaded_ray_installer._install_ray_on_single_node(instance)
+        instances, _ = self.instance_storage.get_instances(
+            instance_ids={instance.instance_id}
+        )
+        assert instances[instance.instance_id].status == Instance.RAY_INSTALLING
+        assert instances[instance.instance_id].version == verison
+
+    @patch.object(RayInstaller, "install_ray")
+    def test_install_ray_on_new_node_install_failed(self, mock_method):
+        self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
+        instance = Instance(
+            instance_id="0",
+            instance_type="worker_nodes1",
+            cloud_instance_id="0",
+            status=Instance.RAY_INSTALLING,
+            node_kind=NodeKind.WORKER,
+        )
+        success, verison = self.instance_storage.upsert_instance(instance)
+        assert success
+        instance.version = verison
+
+        mock_method.side_effect = RuntimeError("Installation failed")
+        self.threaded_ray_installer._install_retry_interval = 0
+        self.threaded_ray_installer._max_install_attempts = 1
+        self.threaded_ray_installer._install_ray_on_single_node(instance)
+
+        instances, _ = self.instance_storage.get_instances(
+            instance_ids={instance.instance_id}
+        )
+        # Make sure the instance status is not updated by the ThreadedRayInstaller
+        # since it should be updated by the Reconciler.
+        assert instances[instance.instance_id].status == Instance.RAY_INSTALLING
+        # Make sure the error is added to the error queue.
+        error = self.error_queue.get()
+        assert error.im_instance_id == instance.instance_id
+        assert "Installation failed" in error.details
+
+    def test_install_ray_on_new_nodes(self):
+        self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
+        instance = Instance(
+            instance_id="0",
+            instance_type="worker_nodes1",
+            cloud_instance_id="0",
+            status=Instance.RAY_INSTALLING,
+            node_kind=NodeKind.WORKER,
+        )
+        success, verison = self.instance_storage.upsert_instance(instance)
+        assert success
+        instance.version = verison
+        self.runner.respond_to_call("json .Config.Env", ["[]" for i in range(1)])
+
+        self.threaded_ray_installer._install_ray_on_new_nodes(instance.instance_id)
+        self.threaded_ray_installer._ray_installation_executor.shutdown(wait=True)
+        instances, _ = self.instance_storage.get_instances(
+            instance_ids={instance.instance_id}
+        )
+        # Make sure the instance status is not updated by the ThreadedRayInstaller
+        # since it should be updated by the Reconciler.
+        assert instances[instance.instance_id].status == Instance.RAY_INSTALLING
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/test_utils.py

@@ -0,0 +1,607 @@
+# coding: utf-8
+import os
+import sys
+from typing import Dict
+
+import pytest  # noqa
+from google.protobuf.json_format import ParseDict
+
+from ray.autoscaler.v2.schema import (
+    ClusterConstraintDemand,
+    ClusterStatus,
+    LaunchRequest,
+    NodeInfo,
+    NodeUsage,
+    PlacementGroupResourceDemand,
+    RayTaskActorDemand,
+    ResourceDemandSummary,
+    ResourceRequestByCount,
+    ResourceUsage,
+    Stats,
+)
+from ray.autoscaler.v2.utils import (
+    ClusterStatusFormatter,
+    ClusterStatusParser,
+    ResourceRequestUtil,
+)
+from ray.core.generated.autoscaler_pb2 import GetClusterStatusReply
+
+
+def _gen_cluster_status_reply(data: Dict):
+    return ParseDict(data, GetClusterStatusReply())
+
+
+class TestResourceRequestUtil:
+    @staticmethod
+    def test_combine_requests_with_affinity():
+
+        AFFINITY = ResourceRequestUtil.PlacementConstraintType.AFFINITY
+        ANTI_AFFINITY = ResourceRequestUtil.PlacementConstraintType.ANTI_AFFINITY
+
+        rqs = [
+            ResourceRequestUtil.make({"CPU": 1}, [(AFFINITY, "1", "1")]),  # 1
+            ResourceRequestUtil.make({"CPU": 2}, [(AFFINITY, "1", "1")]),  # 1
+            ResourceRequestUtil.make({"CPU": 1}, [(AFFINITY, "2", "2")]),  # 2
+            ResourceRequestUtil.make({"CPU": 1}, [(AFFINITY, "2", "2")]),  # 2
+            ResourceRequestUtil.make({"CPU": 1}, [(ANTI_AFFINITY, "2", "2")]),  # 3
+            ResourceRequestUtil.make({"CPU": 1}, [(ANTI_AFFINITY, "2", "2")]),  # 4
+            ResourceRequestUtil.make({"CPU": 1}),  # 5
+        ]
+
+        rq_result = ResourceRequestUtil.combine_requests_with_affinity(rqs)
+        assert len(rq_result) == 5
+        actual = ResourceRequestUtil.to_dict_list(rq_result)
+        expected = [
+            ResourceRequestUtil.to_dict(
+                ResourceRequestUtil.make(
+                    {"CPU": 3},  # Combined
+                    [
+                        (AFFINITY, "1", "1"),
+                    ],
+                )
+            ),
+            ResourceRequestUtil.to_dict(
+                ResourceRequestUtil.make(
+                    {"CPU": 2},  # Combined
+                    [
+                        (AFFINITY, "2", "2"),
+                    ],
+                )
+            ),
+            ResourceRequestUtil.to_dict(
+                ResourceRequestUtil.make(
+                    {"CPU": 1},
+                    [(ANTI_AFFINITY, "2", "2")],
+                )
+            ),
+            ResourceRequestUtil.to_dict(
+                ResourceRequestUtil.make(
+                    {"CPU": 1},
+                    [(ANTI_AFFINITY, "2", "2")],
+                )
+            ),
+            ResourceRequestUtil.to_dict(
+                ResourceRequestUtil.make(
+                    {"CPU": 1},
+                )
+            ),
+        ]
+
+        actual_str_serialized = [str(x) for x in actual]
+        expected_str_serialized = [str(x) for x in expected]
+
+        assert sorted(actual_str_serialized) == sorted(expected_str_serialized)
+
+
+def test_cluster_status_parser_cluster_resource_state():
+    test_data = {
+        "cluster_resource_state": {
+            "node_states": [
+                {
+                    "node_id": b"1" * 4,
+                    "instance_id": "instance1",
+                    "ray_node_type_name": "head_node",
+                    "available_resources": {
+                        "CPU": 0.5,
+                        "GPU": 2.0,
+                    },
+                    "total_resources": {
+                        "CPU": 1,
+                        "GPU": 2.0,
+                    },
+                    "status": "RUNNING",
+                    "node_ip_address": "10.10.10.10",
+                    "instance_type_name": "m5.large",
+                },
+                {
+                    "node_id": b"2" * 4,
+                    "instance_id": "instance2",
+                    "ray_node_type_name": "worker_node",
+                    "available_resources": {},
+                    "total_resources": {
+                        "CPU": 1,
+                        "GPU": 2.0,
+                    },
+                    "status": "DEAD",
+                    "node_ip_address": "22.22.22.22",
+                    "instance_type_name": "m5.large",
+                },
+                {
+                    "node_id": b"3" * 4,
+                    "instance_id": "instance3",
+                    "ray_node_type_name": "worker_node",
+                    "available_resources": {
+                        "CPU": 1.0,
+                        "GPU": 2.0,
+                    },
+                    "total_resources": {
+                        "CPU": 1,
+                        "GPU": 2.0,
+                    },
+                    "idle_duration_ms": 100,
+                    "status": "IDLE",
+                    "node_ip_address": "22.22.22.22",
+                    "instance_type_name": "m5.large",
+                },
+            ],
+            "pending_gang_resource_requests": [
+                {
+                    "requests": [
+                        {
+                            "resources_bundle": {"CPU": 1, "GPU": 1},
+                            "placement_constraints": [
+                                {
+                                    "anti_affinity": {
+                                        "label_name": "_PG_1x1x",
+                                        "label_value": "",
+                                    }
+                                }
+                            ],
+                        },
+                    ],
+                    "details": "1x1x:STRICT_SPREAD|PENDING",
+                },
+                {
+                    "requests": [
+                        {
+                            "resources_bundle": {"GPU": 2},
+                            "placement_constraints": [
+                                {
+                                    "affinity": {
+                                        "label_name": "_PG_2x2x",
+                                        "label_value": "",
+                                    }
+                                }
+                            ],
+                        },
+                    ],
+                    "details": "2x2x:STRICT_PACK|PENDING",
+                },
+            ],
+            "pending_resource_requests": [
+                {
+                    "request": {
+                        "resources_bundle": {"CPU": 1, "GPU": 1},
+                        "placement_constraints": [],
+                    },
+                    "count": 1,
+                },
+            ],
+            "cluster_resource_constraints": [
+                {
+                    "resource_requests": [
+                        {
+                            "request": {
+                                "resources_bundle": {"GPU": 2, "CPU": 100},
+                                "placement_constraints": [],
+                            },
+                            "count": 1,
+                        },
+                    ]
+                }
+            ],
+            "cluster_resource_state_version": 10,
+        },
+        "autoscaling_state": {},
+    }
+    reply = _gen_cluster_status_reply(test_data)
+    stats = Stats(gcs_request_time_s=0.1)
+    cluster_status = ClusterStatusParser.from_get_cluster_status_reply(reply, stats)
+
+    # Assert on health nodes
+    assert len(cluster_status.idle_nodes) + len(cluster_status.active_nodes) == 2
+    assert cluster_status.active_nodes[0].instance_id == "instance1"
+    assert cluster_status.active_nodes[0].ray_node_type_name == "head_node"
+    cluster_status.active_nodes[0].resource_usage.usage.sort(
+        key=lambda x: x.resource_name
+    )
+    assert cluster_status.active_nodes[0].resource_usage == NodeUsage(
+        usage=[
+            ResourceUsage(resource_name="CPU", total=1.0, used=0.5),
+            ResourceUsage(resource_name="GPU", total=2.0, used=0.0),
+        ],
+        idle_time_ms=0,
+    )
+
+    assert cluster_status.idle_nodes[0].instance_id == "instance3"
+    assert cluster_status.idle_nodes[0].ray_node_type_name == "worker_node"
+    cluster_status.idle_nodes[0].resource_usage.usage.sort(
+        key=lambda x: x.resource_name
+    )
+    assert cluster_status.idle_nodes[0].resource_usage == NodeUsage(
+        usage=[
+            ResourceUsage(resource_name="CPU", total=1.0, used=0.0),
+            ResourceUsage(resource_name="GPU", total=2.0, used=0.0),
+        ],
+        idle_time_ms=100,
+    )
+
+    # Assert on dead nodes
+    assert len(cluster_status.failed_nodes) == 1
+    assert cluster_status.failed_nodes[0].instance_id == "instance2"
+    assert cluster_status.failed_nodes[0].ray_node_type_name == "worker_node"
+    assert cluster_status.failed_nodes[0].resource_usage is None
+
+    # Assert on resource demands from tasks
+    assert len(cluster_status.resource_demands.ray_task_actor_demand) == 1
+    assert cluster_status.resource_demands.ray_task_actor_demand[
+        0
+    ].bundles_by_count == [
+        ResourceRequestByCount(
+            bundle={"CPU": 1, "GPU": 1},
+            count=1,
+        )
+    ]
+
+    # Assert on resource demands from placement groups
+    assert len(cluster_status.resource_demands.placement_group_demand) == 2
+    assert sorted(
+        cluster_status.resource_demands.placement_group_demand, key=lambda x: x.pg_id
+    ) == [
+        PlacementGroupResourceDemand(
+            bundles_by_count=[
+                ResourceRequestByCount(bundle={"CPU": 1, "GPU": 1}, count=1)
+            ],
+            strategy="STRICT_SPREAD",
+            pg_id="1x1x",
+            state="PENDING",
+            details="1x1x:STRICT_SPREAD|PENDING",
+        ),
+        PlacementGroupResourceDemand(
+            bundles_by_count=[ResourceRequestByCount(bundle={"GPU": 2}, count=1)],
+            strategy="STRICT_PACK",
+            pg_id="2x2x",
+            state="PENDING",
+            details="2x2x:STRICT_PACK|PENDING",
+        ),
+    ]
+
+    # Assert on resource constraints
+    assert len(cluster_status.resource_demands.cluster_constraint_demand) == 1
+    assert cluster_status.resource_demands.cluster_constraint_demand[
+        0
+    ].bundles_by_count == [
+        ResourceRequestByCount(bundle={"GPU": 2, "CPU": 100}, count=1)
+    ]
+
+    # Assert on the cluster_resource_usage
+    assert sorted(
+        cluster_status.cluster_resource_usage, key=lambda x: x.resource_name
+    ) == [
+        ResourceUsage(resource_name="CPU", total=2.0, used=0.5),
+        ResourceUsage(resource_name="GPU", total=4.0, used=0.0),
+    ]
+
+    # Assert on the node stats
+    assert cluster_status.stats.cluster_resource_state_version == "10"
+    assert cluster_status.stats.gcs_request_time_s == 0.1
+
+
+def test_cluster_status_parser_autoscaler_state():
+    test_data = {
+        "cluster_resource_state": {},
+        "autoscaling_state": {
+            "pending_instance_requests": [
+                {
+                    "instance_type_name": "m5.large",
+                    "ray_node_type_name": "head_node",
+                    "count": 1,
+                    "request_ts": 29999,
+                },
+                {
+                    "instance_type_name": "m5.large",
+                    "ray_node_type_name": "worker_node",
+                    "count": 2,
+                    "request_ts": 19999,
+                },
+            ],
+            "pending_instances": [
+                {
+                    "instance_type_name": "m5.large",
+                    "ray_node_type_name": "head_node",
+                    "instance_id": "instance1",
+                    "ip_address": "10.10.10.10",
+                    "details": "Starting Ray",
+                },
+            ],
+            "failed_instance_requests": [
+                {
+                    "instance_type_name": "m5.large",
+                    "ray_node_type_name": "worker_node",
+                    "count": 2,
+                    "reason": "Insufficient capacity",
+                    "start_ts": 10000,
+                    "failed_ts": 20000,
+                }
+            ],
+            "autoscaler_state_version": 10,
+        },
+    }
+    reply = _gen_cluster_status_reply(test_data)
+    stats = Stats(gcs_request_time_s=0.1)
+    cluster_status = ClusterStatusParser.from_get_cluster_status_reply(reply, stats)
+
+    # Assert on the pending requests
+    assert len(cluster_status.pending_launches) == 2
+    assert cluster_status.pending_launches[0].instance_type_name == "m5.large"
+    assert cluster_status.pending_launches[0].ray_node_type_name == "head_node"
+    assert cluster_status.pending_launches[0].count == 1
+    assert cluster_status.pending_launches[0].request_ts_s == 29999
+    assert cluster_status.pending_launches[1].instance_type_name == "m5.large"
+    assert cluster_status.pending_launches[1].ray_node_type_name == "worker_node"
+    assert cluster_status.pending_launches[1].count == 2
+    assert cluster_status.pending_launches[1].request_ts_s == 19999
+
+    # Assert on the failed requests
+    assert len(cluster_status.failed_launches) == 1
+    assert cluster_status.failed_launches[0].instance_type_name == "m5.large"
+    assert cluster_status.failed_launches[0].ray_node_type_name == "worker_node"
+    assert cluster_status.failed_launches[0].count == 2
+    assert cluster_status.failed_launches[0].details == "Insufficient capacity"
+    assert cluster_status.failed_launches[0].request_ts_s == 10000
+    assert cluster_status.failed_launches[0].failed_ts_s == 20000
+
+    # Assert on the pending nodes
+    assert len(cluster_status.pending_nodes) == 1
+    assert cluster_status.pending_nodes[0].instance_type_name == "m5.large"
+    assert cluster_status.pending_nodes[0].ray_node_type_name == "head_node"
+    assert cluster_status.pending_nodes[0].instance_id == "instance1"
+    assert cluster_status.pending_nodes[0].ip_address == "10.10.10.10"
+    assert cluster_status.pending_nodes[0].details == "Starting Ray"
+
+    # Assert on stats
+    assert cluster_status.stats.autoscaler_version == "10"
+    assert cluster_status.stats.gcs_request_time_s == 0.1
+
+
+def test_cluster_status_formatter():
+    state = ClusterStatus(
+        idle_nodes=[
+            NodeInfo(
+                instance_id="instance1",
+                instance_type_name="m5.large",
+                ray_node_type_name="head_node",
+                ip_address="127.0.0.1",
+                node_status="RUNNING",
+                node_id="fffffffffffffffffffffffffffffffffffffffffffffffffff00001",
+                resource_usage=NodeUsage(
+                    usage=[
+                        ResourceUsage(resource_name="CPU", total=1.0, used=0.5),
+                        ResourceUsage(resource_name="GPU", total=2.0, used=0.0),
+                        ResourceUsage(
+                            resource_name="object_store_memory",
+                            total=10282.0,
+                            used=5555.0,
+                        ),
+                    ],
+                    idle_time_ms=0,
+                ),
+            ),
+            NodeInfo(
+                instance_id="instance2",
+                instance_type_name="m5.large",
+                ray_node_type_name="worker_node",
+                ip_address="127.0.0.2",
+                node_status="RUNNING",
+                node_id="fffffffffffffffffffffffffffffffffffffffffffffffffff00002",
+                resource_usage=NodeUsage(
+                    usage=[
+                        ResourceUsage(resource_name="CPU", total=1.0, used=0),
+                        ResourceUsage(resource_name="GPU", total=2.0, used=0),
+                    ],
+                    idle_time_ms=0,
+                ),
+            ),
+            NodeInfo(
+                instance_id="instance3",
+                instance_type_name="m5.large",
+                ray_node_type_name="worker_node",
+                ip_address="127.0.0.2",
+                node_status="RUNNING",
+                node_id="fffffffffffffffffffffffffffffffffffffffffffffffffff00003",
+                resource_usage=NodeUsage(
+                    usage=[
+                        ResourceUsage(resource_name="CPU", total=1.0, used=0.0),
+                    ],
+                    idle_time_ms=0,
+                ),
+            ),
+        ],
+        pending_launches=[
+            LaunchRequest(
+                instance_type_name="m5.large",
+                count=2,
+                ray_node_type_name="worker_node",
+                state=LaunchRequest.Status.PENDING,
+                request_ts_s=10000,
+            ),
+            LaunchRequest(
+                instance_type_name="g5n.large",
+                count=1,
+                ray_node_type_name="worker_node_gpu",
+                state=LaunchRequest.Status.PENDING,
+                request_ts_s=20000,
+            ),
+        ],
+        failed_launches=[
+            LaunchRequest(
+                instance_type_name="m5.large",
+                count=2,
+                ray_node_type_name="worker_node",
+                state=LaunchRequest.Status.FAILED,
+                details="Insufficient capacity",
+                request_ts_s=10000,
+                failed_ts_s=20000,
+            ),
+        ],
+        pending_nodes=[
+            NodeInfo(
+                instance_id="instance4",
+                instance_type_name="m5.large",
+                ray_node_type_name="worker_node",
+                ip_address="127.0.0.3",
+                details="Starting Ray",
+            ),
+        ],
+        failed_nodes=[
+            NodeInfo(
+                instance_id="instance5",
+                instance_type_name="m5.large",
+                ray_node_type_name="worker_node",
+                ip_address="127.0.0.5",
+                node_status="DEAD",
+            ),
+        ],
+        cluster_resource_usage=[
+            ResourceUsage(resource_name="CPU", total=3.0, used=0.5),
+            ResourceUsage(resource_name="GPU", total=4.0, used=0.0),
+            ResourceUsage(
+                resource_name="object_store_memory", total=10282.0, used=5555.0
+            ),
+        ],
+        resource_demands=ResourceDemandSummary(
+            placement_group_demand=[
+                PlacementGroupResourceDemand(
+                    pg_id="1x1x",
+                    strategy="STRICT_SPREAD",
+                    state="PENDING",
+                    details="1x1x:STRICT_SPREAD|PENDING",
+                    bundles_by_count=[
+                        ResourceRequestByCount(bundle={"CPU": 1, "GPU": 1}, count=1)
+                    ],
+                ),
+                PlacementGroupResourceDemand(
+                    pg_id="2x2x",
+                    strategy="STRICT_PACK",
+                    state="PENDING",
+                    details="2x2x:STRICT_PACK|PENDING",
+                    bundles_by_count=[
+                        ResourceRequestByCount(bundle={"GPU": 2}, count=1)
+                    ],
+                ),
+                PlacementGroupResourceDemand(
+                    pg_id="3x3x",
+                    strategy="STRICT_PACK",
+                    state="PENDING",
+                    details="3x3x:STRICT_PACK|PENDING",
+                    bundles_by_count=[
+                        ResourceRequestByCount(bundle={"GPU": 2}, count=1)
+                    ],
+                ),
+            ],
+            ray_task_actor_demand=[
+                RayTaskActorDemand(
+                    bundles_by_count=[
+                        ResourceRequestByCount(bundle={"CPU": 1, "GPU": 1}, count=1)
+                    ]
+                ),
+                RayTaskActorDemand(
+                    bundles_by_count=[
+                        ResourceRequestByCount(bundle={"CPU": 1, "GPU": 1}, count=10)
+                    ]
+                ),
+            ],
+            cluster_constraint_demand=[
+                ClusterConstraintDemand(
+                    bundles_by_count=[
+                        ResourceRequestByCount(bundle={"GPU": 2, "CPU": 100}, count=2)
+                    ]
+                ),
+            ],
+        ),
+        stats=Stats(
+            gcs_request_time_s=0.1,
+            none_terminated_node_request_time_s=0.2,
+            autoscaler_iteration_time_s=0.3,
+            autoscaler_version="10",
+            cluster_resource_state_version="20",
+            request_ts_s=775303535,
+        ),
+    )
+    actual = ClusterStatusFormatter.format(state, verbose=True)
+
+    expected = """======== Autoscaler status: 1994-07-27 10:05:35 ========
+GCS request time: 0.100000s
+Node Provider non_terminated_nodes time: 0.200000s
+Autoscaler iteration time: 0.300000s
+
+Node status
+--------------------------------------------------------
+Active:
+ (no active nodes)
+Idle:
+ 1 head_node
+ 2 worker_node
+Pending:
+ worker_node, 1 launching
+ worker_node_gpu, 1 launching
+ instance4: worker_node, starting ray
+Recent failures:
+ worker_node: LaunchFailed (latest_attempt: 02:46:40) - Insufficient capacity
+ worker_node: NodeTerminated (instance_id: instance5)
+
+Resources
+--------------------------------------------------------
+Total Usage:
+ 0.5/3.0 CPU
+ 0.0/4.0 GPU
+ 5.42KiB/10.04KiB object_store_memory
+
+From request_resources:
+ {'GPU': 2, 'CPU': 100}: 2 from request_resources()
+Pending Demands:
+ {'CPU': 1, 'GPU': 1}: 11+ pending tasks/actors
+ {'CPU': 1, 'GPU': 1} * 1 (STRICT_SPREAD): 1+ pending placement groups
+ {'GPU': 2} * 1 (STRICT_PACK): 2+ pending placement groups
+
+Node: instance1 (head_node)
+ Id: fffffffffffffffffffffffffffffffffffffffffffffffffff00001
+ Usage:
+  0.5/1.0 CPU
+  0.0/2.0 GPU
+  5.42KiB/10.04KiB object_store_memory
+ Activity:
+  (no activity)
+
+Node: instance2 (worker_node)
+ Id: fffffffffffffffffffffffffffffffffffffffffffffffffff00002
+ Usage:
+  0/1.0 CPU
+  0/2.0 GPU
+ Activity:
+  (no activity)
+
+Node: instance3 (worker_node)
+ Id: fffffffffffffffffffffffffffffffffffffffffffffffffff00003
+ Usage:
+  0.0/1.0 CPU
+ Activity:
+  (no activity)"""
+    assert actual == expected
+
+
+if __name__ == "__main__":
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))

python/ray/autoscaler/v2/tests/util.py

@@ -0,0 +1,206 @@
+import abc
+import operator
+import time
+from abc import abstractmethod
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple
+
+import ray
+from ray.autoscaler.v2.schema import AutoscalerInstance, ClusterStatus, ResourceUsage
+from ray.autoscaler.v2.sdk import get_cluster_status
+from ray.core.generated import autoscaler_pb2
+from ray.core.generated.instance_manager_pb2 import Instance, NodeKind
+
+
+class MockEventLogger:
+    def __init__(self, logger) -> None:
+        self._logs = defaultdict(list)
+        self._logger = logger
+
+    def info(self, s):
+        self._logger.info(s)
+        self._logs["info"].append(s)
+
+    def warning(self, s):
+        self._logger.warning(s)
+        self._logs["warning"].append(s)
+
+    def error(self, s):
+        self._logger.error(s)
+        self._logs["error"].append(s)
+
+    def debug(self, s):
+        self._logger.debug(s)
+        self._logs["debug"].append(s)
+
+    def get_logs(self, level: str) -> List[str]:
+        return self._logs[level]
+
+
+class MockSubscriber:
+    def __init__(self):
+        self.events = []
+
+    def notify(self, events):
+        self.events.extend(events)
+
+    def clear(self):
+        self.events.clear()
+
+    def events_by_id(self, instance_id):
+        return [e for e in self.events if e.instance_id == instance_id]
+
+
+def make_autoscaler_instance(
+    im_instance: Optional[Instance] = None,
+    ray_node: Optional[autoscaler_pb2.NodeState] = None,
+    cloud_instance_id: Optional[str] = None,
+) -> AutoscalerInstance:
+
+    if cloud_instance_id:
+        if im_instance:
+            im_instance.cloud_instance_id = cloud_instance_id
+        if ray_node:
+            ray_node.instance_id = cloud_instance_id
+
+    return AutoscalerInstance(
+        im_instance=im_instance,
+        ray_node=ray_node,
+        cloud_instance_id=cloud_instance_id,
+    )
+
+
+def get_cluster_resource_state(stub) -> autoscaler_pb2.ClusterResourceState:
+    request = autoscaler_pb2.GetClusterResourceStateRequest(
+        last_seen_cluster_resource_state_version=0
+    )
+    return stub.GetClusterResourceState(request).cluster_resource_state
+
+
+class FakeCounter:
+    def dec(self, *args, **kwargs):
+        pass
+
+
+def create_instance(
+    instance_id,
+    status=Instance.UNKNOWN,
+    instance_type="worker_nodes1",
+    status_times: List[Tuple["Instance.InstanceStatus", int]] = None,
+    launch_request_id="",
+    version=0,
+    cloud_instance_id="",
+    ray_node_id="",
+    node_kind=NodeKind.WORKER,
+):
+
+    if not status_times:
+        status_times = [(status, time.time_ns())]
+
+    return Instance(
+        instance_id=instance_id,
+        status=status,
+        version=version,
+        instance_type=instance_type,
+        launch_request_id=launch_request_id,
+        status_history=[
+            Instance.StatusHistory(instance_status=status, timestamp_ns=ts)
+            for status, ts in status_times
+        ],
+        cloud_instance_id=cloud_instance_id,
+        node_id=ray_node_id,
+        node_kind=node_kind,
+    )
+
+
+def report_autoscaling_state(stub, autoscaling_state: autoscaler_pb2.AutoscalingState):
+    request = autoscaler_pb2.ReportAutoscalingStateRequest(
+        autoscaling_state=autoscaling_state
+    )
+    stub.ReportAutoscalingState(request)
+
+
+def get_total_resources(usages: List[ResourceUsage]) -> Dict[str, float]:
+    """Returns a map of resource name to total resource."""
+    return {r.resource_name: r.total for r in usages}
+
+
+def get_available_resources(usages: List[ResourceUsage]) -> Dict[str, float]:
+    """Returns a map of resource name to available resource."""
+    return {r.resource_name: r.total - r.used for r in usages}
+
+
+def get_used_resources(usages: List[ResourceUsage]) -> Dict[str, float]:
+    """Returns a map of resource name to used resource."""
+    return {r.resource_name: r.used for r in usages}
+
+
+"""
+Test utils for e2e autoscaling states checking.
+"""
+
+
+class Check(abc.ABC):
+    @abstractmethod
+    def check(self, status: ClusterStatus):
+        pass
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+
+class CheckFailure(RuntimeError):
+    pass
+
+
+class NodeCountCheck(Check):
+    def __init__(self, count: int):
+        self.count = count
+
+    def check(self, status: ClusterStatus):
+        healthy_nodes = len(status.active_nodes) + len(status.idle_nodes)
+        if healthy_nodes != self.count:
+            raise CheckFailure(f"Expected {self.count} nodes, got {healthy_nodes}")
+
+    def __str__(self) -> str:
+        return f"NodeCountCheck: {self.count}"
+
+
+class TotalResourceCheck(Check):
+    def __init__(
+        self, resources: Dict[str, float], op: operator = operator.eq, enforce_all=False
+    ):
+        self.resources = resources
+        self.op = op
+        self.enforce_all = enforce_all
+
+    def check(self, status: ClusterStatus):
+        actual = status.total_resources()
+        if self.enforce_all and len(actual) != len(self.resources):
+            raise CheckFailure(
+                f"Expected {len(self.resources)} resources, got {len(actual)}"
+            )
+
+        for k, v in self.resources.items():
+            if k not in actual and v:
+                raise CheckFailure(f"Expected resource {k} not found")
+
+            if not self.op(v, actual.get(k, 0)):
+                raise CheckFailure(
+                    f"Expected resource {k} {self.op} {v}, got {actual.get(k, 0)}"
+                )
+
+    def __str__(self) -> str:
+        return f"TotalResourceCheck({self.op}): {self.resources}"
+
+
+def check_cluster(
+    targets: List[Check],
+) -> bool:
+    gcs_address = ray.get_runtime_context().gcs_address
+    cluster_status = get_cluster_status(gcs_address)
+
+    for target in targets:
+        target.check(cluster_status)
+
+    return True