"""Unit test for BatchingNodeProvider. Validates BatchingNodeProvider's book-keeping logic. """ import random import sys from collections import defaultdict from copy import copy from typing import Any, Dict, List, Tuple from uuid import uuid4 import pytest from ray.autoscaler._private.constants import ( DISABLE_LAUNCH_CONFIG_CHECK_KEY, DISABLE_NODE_UPDATERS_KEY, FOREGROUND_NODE_LAUNCH_KEY, ) from ray.autoscaler._private.util import NodeID, NodeType from ray.autoscaler.batching_node_provider import ( BatchingNodeProvider, NodeData, ScaleRequest, ) from ray.autoscaler.tags import ( NODE_KIND_HEAD, NODE_KIND_WORKER, STATUS_UP_TO_DATE, TAG_RAY_NODE_KIND, TAG_RAY_NODE_STATUS, TAG_RAY_REPLICA_INDEX, TAG_RAY_USER_NODE_TYPE, ) class MockBatchingNodeProvider(BatchingNodeProvider): """Mock implementation of a BatchingNodeProvider.""" def __init__( self, provider_config: Dict[str, Any], cluster_name: str, ) -> None: BatchingNodeProvider.__init__(self, provider_config, cluster_name) # Fake cluster manager state: self._node_data_dict: Dict[NodeID, NodeData] = {} self._add_node(node_type="head", node_kind=NODE_KIND_HEAD) # Allow unit test to the control output of safe_to_scale. self._safe_to_scale_flag = True self._scale_request_submitted_count = 0 # Track non_terminated_nodes_calls for use in test_autoscaler.py self.num_non_terminated_nodes_calls = 0 def get_node_data(self) -> Dict[NodeID, NodeData]: self.num_non_terminated_nodes_calls += 1 return self._node_data_dict def set_node_replica_index(self, node_id, replica_index): self._node_data_dict[node_id].replica_index = replica_index def submit_scale_request(self, scale_request: ScaleRequest) -> None: """Simulate modification of cluster state by an external cluster manager.""" self._scale_request_submitted_count += 1 # Delete workers. for node_id in self.scale_request.workers_to_delete: del self._node_data_dict[node_id] # Get counts of workers after the deletion. cur_num_workers = self._cur_num_workers(self._node_data_dict) for node_type in self.scale_request.desired_num_workers: # How many nodes to add. diff = ( self.scale_request.desired_num_workers[node_type] - cur_num_workers[node_type] ) # The next assertion validates the node provider and the test structure. # After removing nodes to terminate, there should be no more # nodes to terminate! assert diff >= 0, diff for _ in range(diff): self._add_node(node_type, NODE_KIND_WORKER) def _add_node(self, node_type, node_kind): new_node_id = str(uuid4()) self._node_data_dict[new_node_id] = NodeData( kind=node_kind, ip=str(uuid4()), status=STATUS_UP_TO_DATE, type=node_type ) def non_terminated_node_ips(self, tag_filters): """This method is used in test_autoscaler.py.""" return [ node_data.ip for node_id, node_data in self._node_data_dict.items() if tag_filters.items() <= self.node_tags(node_id).items() ] def safe_to_scale(self) -> bool: return self.safe_to_scale_flag def _assert_worker_counts( self, expected_worker_counts: Dict[NodeType, int] ) -> None: assert self._cur_num_workers(self._node_data_dict) == expected_worker_counts class BatchingNodeProviderTester: """Utility to test BatchingNodeProvider.""" def __init__(self): self.node_provider = MockBatchingNodeProvider( provider_config={ DISABLE_LAUNCH_CONFIG_CHECK_KEY: True, DISABLE_NODE_UPDATERS_KEY: True, FOREGROUND_NODE_LAUNCH_KEY: True, }, cluster_name="test-cluster", ) # Maps node types to expected node counts. self.expected_node_counts = defaultdict(int) self.expected_node_counts["head"] = 1 # Tracks how many times we expect a scale request to have been submitted. self.expected_scale_request_submitted_count = 0 def update( self, create_node_requests: List[Tuple[str, int]], terminate_nodes_requests: List[Tuple[str, int]], safe_to_scale_flag: bool, ) -> None: """Simulates an autoscaler update with multiple terminate and create calls. Calls non_terminated_nodes, then create/terminate nodes, then post_process. Args: create_node_requests: List of pairs (node type, count). Each pair is used in a create_node call that creates count nodes of the node type. terminate_nodes_requests: List of pairs (node type, count). Each pair is used in a terminate_nodes call that terminates up to count nodes of the node type. safe_to_scale_flag: Passed to the node provider to determine where provider.safe_to_scale() evaluates to True or False. """ self.node_provider.safe_to_scale_flag = safe_to_scale_flag # Call non_terminated_nodes to refresh internal caches. # Also validate node provider state. self.validate_non_terminated_nodes() # Abort if it's not safe to scale. # This behavior is tested in the context of an actual autoscaler update in # test_autoscaler:test_safe_to_scale. if not self.node_provider.safe_to_scale(): return # Terminate some nodes. # Set to track nodes marked for termination during the update. to_terminate_this_update = set() for node_type, count in terminate_nodes_requests: # Each iteration makes a provider.terminate_nodes call. # Terminate "count" nodes of the given node_type. # If count is greater than the number of remaining nodes of the type, # terminate all of the nodes of the type. to_terminate_this_request = [] for node in self.node_provider._node_data_dict: if len(to_terminate_this_request) >= count: break if ( self.node_provider.node_tags(node)[TAG_RAY_USER_NODE_TYPE] != node_type ): continue if node in to_terminate_this_update: continue to_terminate_this_update.add(node) to_terminate_this_request.append(node) self.node_provider.terminate_nodes(to_terminate_this_request) self.expected_node_counts[node_type] -= len(to_terminate_this_request) # else: the scale request will not be submitted. # Create some nodes. for node_type, count in create_node_requests: self.node_provider.create_node( node_config={}, tags={TAG_RAY_USER_NODE_TYPE: node_type}, count=count ) self.expected_node_counts[node_type] += count # else: the scale request will not be submitted. # Scale change is needed exactly when there's something to create or terminate. assert self.node_provider.scale_change_needed is bool( create_node_requests or terminate_nodes_requests ) # Submit the scale request. self.node_provider.post_process() # Expect a scale request to be submitted iff we called create or terminate # at least one. if create_node_requests or terminate_nodes_requests: self.expected_scale_request_submitted_count += 1 def validate_non_terminated_nodes(self): """Calls non_terminated_nodes and validates output against this test classes's accumulated expected state. Tests methods internal_ip, node_tags, non_terminated_nodes of BatchingNodeProvider. """ nodes = self.node_provider.non_terminated_nodes({}) # Get counts of nodes using provider.node_tags to extract node type. actual_node_counts = defaultdict(int) for node in nodes: # Trivial check. Just confirming we can call internal_ip with no issue. assert isinstance(self.node_provider.internal_ip(node), str) # Check tag structure. tags = self.node_provider.node_tags(node) assert set(tags.keys()) == { TAG_RAY_USER_NODE_TYPE, TAG_RAY_NODE_STATUS, TAG_RAY_NODE_KIND, } node_type = tags[TAG_RAY_USER_NODE_TYPE] node_kind = tags[TAG_RAY_NODE_KIND] node_status = tags[TAG_RAY_NODE_STATUS] if node_type == "head": assert node_kind == NODE_KIND_HEAD else: assert node_kind == NODE_KIND_WORKER # Just by construction of this test: assert node_status == STATUS_UP_TO_DATE actual_node_counts[node_type] += 1 # Remove 0 values from expected_node_counts before comparing. for k, v in copy(self.expected_node_counts).items(): if v == 0: del self.expected_node_counts[k] assert actual_node_counts == self.expected_node_counts # Get node counts again using tag filters. actual_node_counts_again = {} for node_type in actual_node_counts: actual_node_counts_again[node_type] = len( self.node_provider.non_terminated_nodes( tag_filters={TAG_RAY_USER_NODE_TYPE: node_type} ) ) assert actual_node_counts_again == self.expected_node_counts # Check filtering by node kind. workers = self.node_provider.non_terminated_nodes( tag_filters={TAG_RAY_NODE_KIND: NODE_KIND_WORKER} ) heads = self.node_provider.non_terminated_nodes( tag_filters={TAG_RAY_NODE_KIND: NODE_KIND_HEAD} ) assert len(heads) == 1 assert set(nodes) == set(workers) | set(heads) # Check filtering by status. up_to_date_nodes = self.node_provider.non_terminated_nodes( tag_filters={TAG_RAY_NODE_STATUS: STATUS_UP_TO_DATE} ) assert set(up_to_date_nodes) == set(nodes) # Make some assertions about internal structure of the node provider. expected_node_counts_without_head = copy(self.expected_node_counts) del expected_node_counts_without_head["head"] # Desired number of workers immediately after calling non_terminated_nodes is # current number of workers. assert ( self.node_provider.scale_request.desired_num_workers == expected_node_counts_without_head ) # scale_change_needed should be reset after calling non_terminated_nodes # (meaning: we've just obtained cluster state and have no indication # from create_node or terminate_node calls that scale needs to change.) assert self.node_provider.scale_change_needed is False # We've submitted the expected number of scale requests: assert ( self.node_provider._scale_request_submitted_count == self.expected_scale_request_submitted_count ) def update_with_random_requests(self): random_requests = self.generate_random_requests() self.update(*random_requests) def generate_random_requests(self): """Generates random sequences of create_node and terminate_nodes requests for the node provider. Generates random safe_to_scale_flag. """ num_creates = random.choice(range(100)) num_terminates = random.choice(range(100)) create_node_requests = [] for _ in range(num_creates): # Choose from 5 worker types. node_type = random.choice([f"type-{x}" for x in range(5)]) # Create up to 9 workers. count = random.choice(range(10)) create_node_requests.append((node_type, count)) terminate_nodes_requests = [] for _ in range(num_terminates): node_type = random.choice([f"type-{x}" for x in range(5)]) # Terminate up to 9 workers. count = random.choice(range(10)) terminate_nodes_requests.append((node_type, count)) # 50% chance of the update being executed. safe_to_scale_flag = random.choice([True, False]) return create_node_requests, terminate_nodes_requests, safe_to_scale_flag def assert_worker_counts(self, expected_worker_counts): """Validates worker counts against internal node provider state.""" self.node_provider._assert_worker_counts(expected_worker_counts) @pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.") def test_batching_node_provider_basic(): tester = BatchingNodeProviderTester() tester.update( create_node_requests=[ ("type-1", 5), ], terminate_nodes_requests=[], safe_to_scale_flag=True, ) tester.assert_worker_counts({"type-1": 5}) assert tester.node_provider._scale_request_submitted_count == 1 tester.update( create_node_requests=[("type-2", 5), ("type-2", 5)], terminate_nodes_requests=[("type-1", 2)], safe_to_scale_flag=True, ) tester.assert_worker_counts({"type-1": 3, "type-2": 10}) assert tester.node_provider._scale_request_submitted_count == 2 tester.update( create_node_requests=[], terminate_nodes_requests=[("type-1", 2), ("type-2", 1), ("type-2", 1)], safe_to_scale_flag=True, ) tester.assert_worker_counts({"type-1": 1, "type-2": 8}) assert tester.node_provider._scale_request_submitted_count == 3 tester.update( create_node_requests=[], terminate_nodes_requests=[], safe_to_scale_flag=True, ) tester.assert_worker_counts({"type-1": 1, "type-2": 8}) # No scale request submitted, since there were no create/terminate calls. assert tester.node_provider._scale_request_submitted_count == 3 @pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.") def test_batching_node_provider_many_requests(): """Simulate 10 autoscaler updates with randomly generated create/terminate requests. """ tester = BatchingNodeProviderTester() for _ in range(2): tester.update_with_random_requests() # Final check. tester.validate_non_terminated_nodes() @pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.") def test_terminate_safeguards(): """Tests the following behaviors: - the node provider ignores requests to terminate a node twice. - the node provider ignores requests to terminate an unknown node. """ node_provider = MockBatchingNodeProvider( provider_config={ DISABLE_LAUNCH_CONFIG_CHECK_KEY: True, DISABLE_NODE_UPDATERS_KEY: True, FOREGROUND_NODE_LAUNCH_KEY: True, }, cluster_name="test-cluster", ) nodes = node_provider.non_terminated_nodes({}) assert len(nodes) == 1 head_node = nodes[0] node_provider.create_node( node_config={}, tags={TAG_RAY_USER_NODE_TYPE: "type"}, count=1 ) node_provider.post_process() nodes = node_provider.non_terminated_nodes({}) assert len(nodes) == 2 worker_node = "" for node in nodes: if node == head_node: continue else: worker_node = node # This node is not in our list. unknown_node = node + worker_node node_provider.terminate_node(unknown_node) node_provider.post_process() nodes = node_provider.non_terminated_nodes({}) # Terminate request was ignored because the node is unknown. assert len(nodes) == 2 node_provider.terminate_node(worker_node) node_provider.terminate_node(worker_node) node_provider.post_process() nodes = node_provider.non_terminated_nodes({}) # Second terminate request was ignored. assert len(nodes) == 1 @pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.") def test_terminate_node_in_multihost_replica(): """Test multi-host replica deletion logic for KubeRay. Tests manually deleting a node in a multi-host replica and verifying that the entire replica is scaled down. Nodes belonging to the same multi-host replica are identified through a replicaIndex label set by a GKE webhook. """ # create 4 TPU nodes with MockBatchingNodeProvider node_provider = MockBatchingNodeProvider( provider_config={ DISABLE_LAUNCH_CONFIG_CHECK_KEY: True, DISABLE_NODE_UPDATERS_KEY: True, FOREGROUND_NODE_LAUNCH_KEY: True, }, cluster_name="test-cluster", ) num_tpu_workers = 4 for i in range(num_tpu_workers): node_provider._add_node(node_type="TPU", node_kind=NODE_KIND_WORKER) # Set replica_index in node_data for all workers workers = node_provider.non_terminated_nodes( tag_filters={TAG_RAY_NODE_KIND: NODE_KIND_WORKER} ) assert len(workers) == num_tpu_workers for index, node_id in enumerate(workers): if index < num_tpu_workers // 2: node_provider.set_node_replica_index(node_id, "tpu-group-0") else: node_provider.set_node_replica_index(node_id, "tpu-group-1") # Verify RAY_REPLICA_INDEX tag has been set replicaIndexFilter = {TAG_RAY_REPLICA_INDEX: "tpu-group-0"} replicaWorkers1 = node_provider.non_terminated_nodes(tag_filters=replicaIndexFilter) assert len(replicaWorkers1) == num_tpu_workers // 2 replicaIndexFilter[TAG_RAY_REPLICA_INDEX] = "tpu-group-1" replicaWorkers2 = node_provider.non_terminated_nodes(tag_filters=replicaIndexFilter) assert len(replicaWorkers2) == num_tpu_workers // 2 # Verify replica_to_nodes mapping has been populated assert ( len(node_provider.replica_index_to_nodes["tpu-group-0"]) == num_tpu_workers // 2 ) assert ( len(node_provider.replica_index_to_nodes["tpu-group-1"]) == num_tpu_workers // 2 ) worker_0 = replicaWorkers1[0] # tpu-group-0 worker_2 = replicaWorkers2[0] # tpu-group-1 # Manually delete one TPU worker in tpu-group-0 # BatchingNodeProvider should scale down all nodes in the replica assert worker_0 in node_provider.node_data_dict node_provider.terminate_node(worker_0) assert len(node_provider.scale_request.workers_to_delete) == num_tpu_workers // 2 # Scale down the tpu-group-1 replica assert worker_2 in node_provider.node_data_dict node_provider.terminate_node(worker_2) assert len(node_provider.scale_request.workers_to_delete) == num_tpu_workers if __name__ == "__main__": sys.exit(pytest.main(["-sv", __file__]))