# 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"])