import copy import sys from collections import defaultdict from pathlib import Path from typing import List, Set from unittest import mock import jsonpatch import pytest import yaml from ray.autoscaler._private.kuberay.node_provider import ( KubeRayNodeProvider, ScaleRequest, _worker_group_index, _worker_group_max_replicas, _worker_group_replicas, ) from ray.autoscaler._private.util import NodeID from ray.autoscaler.batching_node_provider import NodeData from ray.tests.kuberay.test_autoscaling_config import get_basic_ray_cr def _get_basic_ray_cr_workers_to_delete( cpu_workers_to_delete: List[NodeID], gpu_workers_to_delete: List[NodeID], tpu_workers_to_delete: List[NodeID], ): """Generate a Ray cluster with non-empty workersToDelete field.""" raycluster = get_basic_ray_cr() raycluster["spec"]["workerGroupSpecs"][0]["scaleStrategy"] = { "workersToDelete": cpu_workers_to_delete } raycluster["spec"]["workerGroupSpecs"][1]["scaleStrategy"] = { "workersToDelete": gpu_workers_to_delete } raycluster["spec"]["workerGroupSpecs"][2]["scaleStrategy"] = { "workersToDelete": tpu_workers_to_delete } return raycluster def _get_test_yaml(file_name): file_path = str(Path(__file__).resolve().parent / "test_files" / file_name) return yaml.safe_load(open(file_path).read()) @pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.") @pytest.mark.parametrize( "group_name,expected_index", [("small-group", 0), ("gpu-group", 1)] ) def test_worker_group_index(group_name, expected_index): """Basic unit test for _worker_group_index. Uses a RayCluster CR with worker groups "small-group" and "gpu-group", listed in that order. """ raycluster_cr = get_basic_ray_cr() assert _worker_group_index(raycluster_cr, group_name) == expected_index @pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.") @pytest.mark.parametrize( "group_index,expected_max_replicas,expected_replicas", [(0, 300, 1), (1, 200, 1), (2, 4, 1), (3, None, 0)], ) def test_worker_group_replicas(group_index, expected_max_replicas, expected_replicas): """Basic unit test for _worker_group_max_replicas and _worker_group_replicas Uses a RayCluster CR with worker groups with 300 maxReplicas, 200 maxReplicas, and unspecified maxReplicas, in that order. """ raycluster = get_basic_ray_cr() # Add a worker group without maxReplicas to confirm behavior # when maxReplicas is not specified. no_max_replicas_group = copy.deepcopy(raycluster["spec"]["workerGroupSpecs"][0]) no_max_replicas_group["groupName"] = "no-max-replicas" del no_max_replicas_group["maxReplicas"] # Also, replicas field, just for the sake of testing. no_max_replicas_group["replicas"] = 0 raycluster["spec"]["workerGroupSpecs"].append(no_max_replicas_group) assert _worker_group_max_replicas(raycluster, group_index) == expected_max_replicas assert _worker_group_replicas(raycluster, group_index) == expected_replicas @pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.") @pytest.mark.parametrize( "attempted_target_replica_count,expected_target_replica_count", [(200, 200), (250, 250), (300, 300), (400, 300), (1000, 300)], ) def test_create_node_cap_at_max( attempted_target_replica_count: int, expected_target_replica_count: int ): """Validates that KubeRayNodeProvider does not attempt to create more nodes than allowed by maxReplicas. For the config in this test, maxReplicas is fixed at 300. Args: attempted_target_replica_count: The mocked desired replica count for a given worker group. expected_target_replica_count: The actual requested replicaCount. Should be capped at maxReplicas (300, for the config in this test.) """ raycluster = get_basic_ray_cr() with mock.patch.object(KubeRayNodeProvider, "__init__", return_value=None): kr_node_provider = KubeRayNodeProvider(provider_config={}, cluster_name="fake") scale_request = ScaleRequest( workers_to_delete=set(), desired_num_workers={"small-group": attempted_target_replica_count}, ) patch = kr_node_provider._scale_request_to_patch_payload( scale_request=scale_request, raycluster=raycluster ) assert patch[0]["value"] == expected_target_replica_count @pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.") @pytest.mark.parametrize( "podlist_file,expected_node_data", [ ( # Pod list obtained by running kubectl get pod -o yaml at runtime. "podlist1.yaml", { "raycluster-autoscaler-head-8zsc8": NodeData( kind="head", type="headgroup", replica_index=None, ip="10.4.2.6", status="up-to-date", ), # up-to-date status because the Ray container is in running status "raycluster-autoscaler-worker-small-group-dkz2r": NodeData( kind="worker", type="small-group", replica_index=None, ip="10.4.1.8", status="waiting", ), # waiting status, because Ray container's state is "waiting". # The pod list includes a worker with non-null deletion timestamp. # It is excluded from the node data because it is considered # "terminated". }, ), ( # Pod list obtained by running kubectl get pod -o yaml at runtime. "podlist2.yaml", { "raycluster-autoscaler-head-8zsc8": NodeData( kind="head", type="headgroup", replica_index=None, ip="10.4.2.6", status="up-to-date", ), "raycluster-autoscaler-worker-fake-gpu-group-2qnhv": NodeData( kind="worker", type="fake-gpu-group", replica_index=None, ip="10.4.0.6", status="up-to-date", ), "raycluster-autoscaler-worker-small-group-dkz2r": NodeData( kind="worker", type="small-group", replica_index=None, ip="10.4.1.8", status="up-to-date", ), "raycluster-autoscaler-worker-small-group-lbfm4": NodeData( kind="worker", type="small-group", replica_index=None, ip="10.4.0.5", status="up-to-date", ), "raycluster-autoscaler-tpu-group-worker-s8jhq": NodeData( kind="worker", type="tpu-group", replica_index="tpu-group-0", ip="10.24.9.4", status="up-to-date", ), "raycluster-autoscaler-tpu-group-worker-jd69f": NodeData( kind="worker", type="tpu-group", replica_index="tpu-group-0", ip="10.24.8.4", status="up-to-date", ), }, ), ], ) def test_get_node_data(podlist_file: str, expected_node_data): """Test translation of a K8s pod list into autoscaler node data.""" pod_list = _get_test_yaml(podlist_file) def mock_get(node_provider, path): if "pods" in path: return pod_list elif "raycluster" in path: return get_basic_ray_cr() else: raise ValueError("Invalid path.") with mock.patch.object( KubeRayNodeProvider, "__init__", return_value=None ), mock.patch.object(KubeRayNodeProvider, "_get", mock_get): kr_node_provider = KubeRayNodeProvider(provider_config={}, cluster_name="fake") kr_node_provider.cluster_name = "fake" kr_node_provider.replica_index_to_nodes = defaultdict(list[str]) nodes = kr_node_provider.non_terminated_nodes({}) assert kr_node_provider.node_data_dict == expected_node_data assert set(nodes) == set(expected_node_data.keys()) @pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.") @pytest.mark.parametrize( "node_data_dict,scale_request,expected_patch_payload", [ ( { "raycluster-autoscaler-head-8zsc8": NodeData( kind="head", type="headgroup", replica_index=None, ip="10.4.2.6", status="up-to-date", ), "raycluster-autoscaler-worker-fake-gpu-group-2qnhv": NodeData( kind="worker", type="fake-gpu-group", replica_index=None, ip="10.4.0.6", status="up-to-date", ), "raycluster-autoscaler-worker-small-group-dkz2r": NodeData( kind="worker", type="small-group", replica_index=None, ip="10.4.1.8", status="up-to-date", ), "raycluster-autoscaler-worker-small-group-lbfm4": NodeData( kind="worker", type="small-group", replica_index=None, ip="10.4.0.5", status="up-to-date", ), }, ScaleRequest( desired_num_workers={ "small-group": 1, # Delete 1 "gpu-group": 1, # Don't touch "blah-group": 5, # Create 5 }, workers_to_delete={ "raycluster-autoscaler-worker-small-group-dkz2r", }, ), [ { "op": "replace", "path": "/spec/workerGroupSpecs/3/replicas", "value": 5, }, { "op": "replace", "path": "/spec/workerGroupSpecs/0/scaleStrategy", "value": { "workersToDelete": [ "raycluster-autoscaler-worker-small-group-dkz2r" ] }, }, ], ), ], ) def test_submit_scale_request(node_data_dict, scale_request, expected_patch_payload): """Test the KubeRayNodeProvider's RayCluster patch payload given a dict of current node counts and a scale request. """ raycluster = get_basic_ray_cr() # Add another worker group for the sake of this test. blah_group = copy.deepcopy(raycluster["spec"]["workerGroupSpecs"][1]) blah_group["groupName"] = "blah-group" raycluster["spec"]["workerGroupSpecs"].append(blah_group) with mock.patch.object(KubeRayNodeProvider, "__init__", return_value=None): kr_node_provider = KubeRayNodeProvider(provider_config={}, cluster_name="fake") kr_node_provider.node_data_dict = node_data_dict patch_payload = kr_node_provider._scale_request_to_patch_payload( scale_request=scale_request, raycluster=raycluster ) assert patch_payload == expected_patch_payload @pytest.mark.parametrize("node_set", [{"A", "B", "C", "D", "E"}]) @pytest.mark.parametrize("cpu_workers_to_delete", ["A", "Z"]) @pytest.mark.parametrize("gpu_workers_to_delete", ["B", "Y"]) @pytest.mark.parametrize("tpu_workers_to_delete", ["C", "X"]) @pytest.mark.skipif(sys.platform.startswith("win"), reason="Not relevant on Windows.") def test_safe_to_scale( node_set: Set[NodeID], cpu_workers_to_delete: List[NodeID], gpu_workers_to_delete: List[NodeID], tpu_workers_to_delete: List[NodeID], ): # NodeData values unimportant for this test. mock_node_data = NodeData("-", "-", "-", "-", "-") node_data_dict = {node_id: mock_node_data for node_id in node_set} raycluster = _get_basic_ray_cr_workers_to_delete( cpu_workers_to_delete, gpu_workers_to_delete, tpu_workers_to_delete ) def mock_patch(kuberay_provider, path, patch_payload): patch = jsonpatch.JsonPatch(patch_payload) kuberay_provider._patched_raycluster = patch.apply(kuberay_provider._raycluster) with mock.patch.object( KubeRayNodeProvider, "__init__", return_value=None ), mock.patch.object(KubeRayNodeProvider, "_patch", mock_patch): kr_node_provider = KubeRayNodeProvider(provider_config={}, cluster_name="fake") kr_node_provider.cluster_name = "fake" kr_node_provider._patched_raycluster = raycluster kr_node_provider._raycluster = raycluster kr_node_provider.node_data_dict = node_data_dict actual_safe = kr_node_provider.safe_to_scale() expected_safe = ( not any( cpu_worker_to_delete in node_set for cpu_worker_to_delete in cpu_workers_to_delete ) and not any( gpu_worker_to_delete in node_set for gpu_worker_to_delete in gpu_workers_to_delete ) and not any( tpu_worker_to_delete in node_set for tpu_worker_to_delete in tpu_workers_to_delete ) ) assert expected_safe is actual_safe patched_cpu_workers_to_delete = kr_node_provider._patched_raycluster["spec"][ "workerGroupSpecs" ][0]["scaleStrategy"]["workersToDelete"] patched_gpu_workers_to_delete = kr_node_provider._patched_raycluster["spec"][ "workerGroupSpecs" ][1]["scaleStrategy"]["workersToDelete"] patched_tpu_workers_to_delete = kr_node_provider._patched_raycluster["spec"][ "workerGroupSpecs" ][2]["scaleStrategy"]["workersToDelete"] if expected_safe: # Cleaned up workers to delete assert patched_cpu_workers_to_delete == [] assert patched_gpu_workers_to_delete == [] assert patched_tpu_workers_to_delete == [] else: # Did not clean up workers to delete assert patched_cpu_workers_to_delete == cpu_workers_to_delete assert patched_gpu_workers_to_delete == gpu_workers_to_delete assert patched_tpu_workers_to_delete == tpu_workers_to_delete if __name__ == "__main__": sys.exit(pytest.main(["-v", __file__]))