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__]))