wehub-resource-sync
1005 lines
36 KiB
Python
1005 lines
36 KiB
Python
import logging
|
|
import os
|
|
import sys
|
|
import time
|
|
import unittest
|
|
|
|
# coding: utf-8
|
|
# coding: utf-8
|
|
from collections import defaultdict
|
|
from typing import Any, Dict, List, Union
|
|
from unittest import mock
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest # noqa
|
|
|
|
import ray
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray._private.test_utils import get_test_config_path
|
|
from ray.autoscaler._private.constants import (
|
|
AUTOSCALER_MAX_CONCURRENT_LAUNCHES,
|
|
AUTOSCALER_MAX_LAUNCH_BATCH,
|
|
)
|
|
from ray.autoscaler._private.fake_multi_node.node_provider import FakeMultiNodeProvider
|
|
from ray.autoscaler._private.kuberay.node_provider import IKubernetesHttpApiClient
|
|
from ray.autoscaler.v2.instance_manager.cloud_providers.kuberay.cloud_provider import (
|
|
NO_DRIVER_TTL_EXPIRED_ANNOTATION,
|
|
KubeRayProvider,
|
|
)
|
|
from ray.autoscaler.v2.instance_manager.config import (
|
|
AutoscalingConfig,
|
|
FileConfigReader,
|
|
)
|
|
from ray.autoscaler.v2.instance_manager.node_provider import (
|
|
CloudInstance,
|
|
ICloudInstanceProvider,
|
|
LaunchNodeError,
|
|
NodeProviderAdapter,
|
|
TerminateNodeError,
|
|
logger,
|
|
)
|
|
from ray.core.generated.instance_manager_pb2 import NodeKind
|
|
from ray.tests.autoscaler_test_utils import MockProvider
|
|
from ray.tests.kuberay.test_autoscaling_config import get_basic_ray_cr
|
|
from ray.tests.kuberay.test_kuberay_node_provider import _get_test_yaml
|
|
|
|
logger.setLevel(logging.DEBUG)
|
|
|
|
|
|
class CloudInstanceProviderTesterBase(ICloudInstanceProvider):
|
|
def __init__(
|
|
self,
|
|
inner_provider: ICloudInstanceProvider,
|
|
config: AutoscalingConfig,
|
|
):
|
|
self.inner_provider = inner_provider
|
|
self.config = config
|
|
|
|
def __del__(self):
|
|
self.shutdown()
|
|
|
|
def shutdown(self):
|
|
pass
|
|
|
|
def launch(self, request_id, shape):
|
|
self.inner_provider.launch(shape=shape, request_id=request_id)
|
|
|
|
def terminate(self, request_id, ids):
|
|
self.inner_provider.terminate(ids=ids, request_id=request_id)
|
|
|
|
def poll_errors(self):
|
|
return self.inner_provider.poll_errors()
|
|
|
|
def get_non_terminated(self):
|
|
return self.inner_provider.get_non_terminated()
|
|
|
|
############################
|
|
# Test mock methods
|
|
############################
|
|
def _add_creation_error(self, e: Exception):
|
|
raise NotImplementedError("Subclass should implement it")
|
|
|
|
def _add_termination_errors(self, e: Exception):
|
|
raise NotImplementedError("Subclass should implement it")
|
|
|
|
|
|
class FakeMultiNodeProviderTester(CloudInstanceProviderTesterBase):
|
|
def __init__(self, **kwargs):
|
|
self.config_reader = FileConfigReader(
|
|
get_test_config_path("test_ray_complex.yaml"), skip_content_hash=True
|
|
)
|
|
self.config = self.config_reader.get_cached_autoscaling_config()
|
|
self.ray_session = None
|
|
|
|
os.environ["RAY_FAKE_CLUSTER"] = "1"
|
|
provider_config = self.config.get_provider_config()
|
|
# This is a bit hacky but we need a fake head node.
|
|
self.ray_session = ray.init()
|
|
provider_config["gcs_address"] = self.ray_session.address_info["gcs_address"]
|
|
provider_config["head_node_id"] = self.ray_session.address_info["node_id"]
|
|
provider_config["launch_multiple"] = True
|
|
self.base_provider = FakeMultiNodeProvider(
|
|
provider_config,
|
|
cluster_name="test",
|
|
)
|
|
|
|
provider = NodeProviderAdapter(
|
|
self.base_provider,
|
|
self.config_reader,
|
|
)
|
|
super().__init__(provider, self.config)
|
|
|
|
def get_non_terminated(self):
|
|
nodes = self.inner_provider.get_non_terminated()
|
|
nodes.pop(self.ray_session.address_info["node_id"], None)
|
|
return nodes
|
|
|
|
def shutdown(self):
|
|
ray.shutdown()
|
|
|
|
def _add_creation_error(self, e: Exception):
|
|
self.base_provider._test_set_creation_error(e)
|
|
|
|
def _add_termination_errors(self, e: Exception):
|
|
self.base_provider._test_add_termination_errors(e)
|
|
|
|
|
|
class MockProviderTester(CloudInstanceProviderTesterBase):
|
|
def __init__(self, **kwargs):
|
|
self.config_reader = FileConfigReader(
|
|
get_test_config_path("test_ray_complex.yaml"), skip_content_hash=True
|
|
)
|
|
self.config = self.config_reader.get_cached_autoscaling_config()
|
|
self.base_provider = MockProvider()
|
|
provider = NodeProviderAdapter(
|
|
self.base_provider,
|
|
self.config_reader,
|
|
)
|
|
super().__init__(provider, self.config)
|
|
|
|
def _add_creation_error(self, e: Exception):
|
|
self.base_provider.creation_error = e
|
|
|
|
def _add_termination_errors(self, e: Exception):
|
|
self.base_provider.termination_errors = e
|
|
|
|
|
|
class MagicMockProviderTester(CloudInstanceProviderTesterBase):
|
|
def __init__(
|
|
self,
|
|
max_concurrent_launches=AUTOSCALER_MAX_CONCURRENT_LAUNCHES,
|
|
max_launch_batch_per_type=AUTOSCALER_MAX_LAUNCH_BATCH,
|
|
**kwargs,
|
|
):
|
|
self.config_reader = FileConfigReader(
|
|
get_test_config_path("test_ray_complex.yaml"), skip_content_hash=True
|
|
)
|
|
self.config = self.config_reader.get_cached_autoscaling_config()
|
|
self.base_provider = MagicMock()
|
|
provider = NodeProviderAdapter(
|
|
self.base_provider,
|
|
self.config_reader,
|
|
max_launch_batch_per_type=max_launch_batch_per_type,
|
|
max_concurrent_launches=max_concurrent_launches,
|
|
)
|
|
super().__init__(provider, self.config)
|
|
|
|
def _add_creation_error(self, e: Exception):
|
|
self.base_provider.create_node_with_resources_and_labels.side_effect = e
|
|
|
|
def _add_termination_errors(self, e: Exception):
|
|
self.base_provider.terminate_nodes.side_effect = e
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def get_provider():
|
|
def _get_provider(name, **kwargs):
|
|
if name == "fake_multi":
|
|
provider = FakeMultiNodeProviderTester(**kwargs)
|
|
elif name == "mock":
|
|
provider = MockProviderTester(**kwargs)
|
|
elif name == "magic_mock":
|
|
provider = MagicMockProviderTester(**kwargs)
|
|
else:
|
|
raise ValueError(f"Invalid provider type: {name}")
|
|
|
|
return provider
|
|
|
|
yield _get_provider
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"provider_name",
|
|
["fake_multi", "mock"],
|
|
)
|
|
def test_node_providers_basic(get_provider, provider_name):
|
|
# Test launching.
|
|
provider = get_provider(name=provider_name)
|
|
timeout_s = 30 if provider_name == "fake_multi" else 10
|
|
provider.launch(
|
|
shape={"worker_nodes": 2},
|
|
request_id="1",
|
|
)
|
|
|
|
provider.launch(
|
|
request_id="2",
|
|
shape={"worker_nodes": 2, "worker_nodes1": 1},
|
|
)
|
|
|
|
def verify():
|
|
nodes_by_type = defaultdict(int)
|
|
for node in provider.get_non_terminated().values():
|
|
nodes_by_type[node.node_type] += 1
|
|
errors = provider.poll_errors()
|
|
print(errors)
|
|
assert nodes_by_type == {"worker_nodes": 4, "worker_nodes1": 1}
|
|
return True
|
|
|
|
wait_for_condition(verify, timeout=timeout_s)
|
|
|
|
nodes = provider.get_non_terminated().keys()
|
|
|
|
# Terminate them all
|
|
provider.terminate(
|
|
ids=nodes,
|
|
request_id="3",
|
|
)
|
|
|
|
# Launch some.
|
|
provider.launch(
|
|
shape={"worker_nodes": 1},
|
|
request_id="4",
|
|
)
|
|
|
|
def verify():
|
|
nodes_by_type = defaultdict(int)
|
|
for node in provider.get_non_terminated().values():
|
|
nodes_by_type[node.node_type] += 1
|
|
|
|
assert nodes_by_type == {"worker_nodes": 1}
|
|
for node in provider.get_non_terminated().values():
|
|
assert node.request_id == "4"
|
|
return True
|
|
|
|
wait_for_condition(verify, timeout=timeout_s)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"provider_name",
|
|
["fake_multi", "mock"],
|
|
)
|
|
def test_launch_failure(get_provider, provider_name):
|
|
provider = get_provider(name=provider_name)
|
|
provider._add_creation_error(Exception("failed to create node"))
|
|
|
|
provider.launch(
|
|
shape={"worker_nodes": 2},
|
|
request_id="2",
|
|
)
|
|
|
|
def verify():
|
|
errors = provider.poll_errors()
|
|
assert len(errors) == 1
|
|
assert isinstance(errors[0], LaunchNodeError)
|
|
assert errors[0].node_type == "worker_nodes"
|
|
assert errors[0].request_id == "2"
|
|
return True
|
|
|
|
wait_for_condition(verify)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"provider_name",
|
|
["fake_multi", "mock"],
|
|
)
|
|
def test_terminate_node_failure(get_provider, provider_name):
|
|
provider = get_provider(name=provider_name)
|
|
provider._add_termination_errors(Exception("failed to terminate node"))
|
|
|
|
provider.launch(request_id="launch1", shape={"worker_nodes": 1})
|
|
|
|
def nodes_launched():
|
|
nodes = provider.get_non_terminated()
|
|
return len(nodes) == 1
|
|
|
|
wait_for_condition(nodes_launched)
|
|
|
|
provider.terminate(request_id="terminate1", ids=["0"])
|
|
|
|
def verify():
|
|
errors = provider.poll_errors()
|
|
nodes = provider.get_non_terminated()
|
|
assert len(nodes) == 1
|
|
assert len(errors) == 1
|
|
assert isinstance(errors[0], TerminateNodeError)
|
|
assert errors[0].cloud_instance_id == "0"
|
|
assert errors[0].request_id == "terminate1"
|
|
return True
|
|
|
|
wait_for_condition(verify)
|
|
|
|
|
|
def test_launch_executor_concurrency(get_provider):
|
|
import threading
|
|
|
|
provider = get_provider(
|
|
name="magic_mock", max_concurrent_launches=1, max_launch_batch_per_type=1
|
|
)
|
|
|
|
launch_event = threading.Event()
|
|
|
|
def loop(*args, **kwargs):
|
|
launch_event.wait()
|
|
|
|
provider.base_provider.create_node_with_resources_and_labels.side_effect = loop
|
|
|
|
provider.launch(
|
|
shape={
|
|
"worker_nodes": 1,
|
|
"worker_nodes1": 1,
|
|
}, # 2 types, but concurrent types to launch is 1.
|
|
request_id="1",
|
|
)
|
|
# Assert called only once.
|
|
for _ in range(10):
|
|
assert (
|
|
provider.base_provider.create_node_with_resources_and_labels.call_count <= 1
|
|
)
|
|
time.sleep(0.1)
|
|
|
|
# Finish the call.
|
|
launch_event.set()
|
|
|
|
def verify():
|
|
assert (
|
|
provider.base_provider.create_node_with_resources_and_labels.call_count == 2
|
|
)
|
|
return True
|
|
|
|
wait_for_condition(verify)
|
|
|
|
|
|
#######################################
|
|
# Integration test for KubeRay Provider
|
|
#######################################
|
|
class MockKubernetesHttpApiClient(IKubernetesHttpApiClient):
|
|
def __init__(self, pod_list: List[Dict[str, Any]], ray_cluster: Dict[str, Any]):
|
|
self._ray_cluster = ray_cluster
|
|
self._pod_list = pod_list
|
|
self._patches = {}
|
|
|
|
def get(self, path: str) -> Dict[str, Any]:
|
|
if "pods" in path:
|
|
return self._pod_list
|
|
if "rayclusters" in path:
|
|
return self._ray_cluster
|
|
|
|
raise NotImplementedError(f"get {path}")
|
|
|
|
def patch(
|
|
self,
|
|
path: str,
|
|
patches: Union[List[Dict[str, Any]], Dict[str, Any]],
|
|
content_type: str = "application/json-patch+json",
|
|
):
|
|
self._patches[path] = patches
|
|
return {path: patches}
|
|
|
|
def get_patches(self, path: str) -> Union[List[Dict[str, Any]], Dict[str, Any]]:
|
|
return self._patches[path]
|
|
|
|
|
|
class KubeRayProviderIntegrationTest(unittest.TestCase):
|
|
def setUp(self):
|
|
raycluster_cr = get_basic_ray_cr()
|
|
# Remove fake TPU and GPU worker groups from CR since podlist1 only
|
|
# contains small-group.
|
|
raycluster_cr["spec"]["workerGroupSpecs"][1]["replicas"] = 0
|
|
raycluster_cr["spec"]["workerGroupSpecs"][2]["replicas"] = 0
|
|
self.mock_client = MockKubernetesHttpApiClient(
|
|
_get_test_yaml("podlist1.yaml"), raycluster_cr
|
|
)
|
|
self.provider = KubeRayProvider(
|
|
cluster_name="test",
|
|
provider_config={
|
|
"namespace": "default",
|
|
"head_node_type": "headgroup",
|
|
},
|
|
gcs_client=MagicMock(),
|
|
k8s_api_client=self.mock_client,
|
|
)
|
|
# In production _sync_with_api_server caches the CR before the
|
|
# no-driver annotation is set; mirror that for the dispatch tests.
|
|
self.provider._ray_cluster = raycluster_cr
|
|
|
|
def test_get_nodes(self):
|
|
nodes = self.provider.get_non_terminated()
|
|
errors = self.provider.poll_errors()
|
|
|
|
assert len(nodes) == 2
|
|
assert len(errors) == 0
|
|
|
|
assert sorted(nodes) == sorted(
|
|
{
|
|
"raycluster-autoscaler-head-8zsc8": CloudInstance(
|
|
cloud_instance_id="raycluster-autoscaler-head-8zsc8",
|
|
node_kind=NodeKind.HEAD,
|
|
node_type="headgroup",
|
|
is_running=True,
|
|
), # up-to-date status because the Ray container is in running status
|
|
"raycluster-autoscaler-worker-small-group-dkz2r": CloudInstance(
|
|
cloud_instance_id="raycluster-autoscaler-worker-small-group-dkz2r",
|
|
node_kind=NodeKind.WORKER,
|
|
node_type="small-group",
|
|
is_running=False,
|
|
), # waiting status, because Ray container's state is pending.
|
|
}
|
|
)
|
|
|
|
def test_launch_node(self):
|
|
launch_request = {"small-group": 1}
|
|
self.provider.launch(shape=launch_request, request_id="launch-1")
|
|
|
|
patches = self.mock_client.get_patches(
|
|
f"rayclusters/{self.provider._cluster_name}"
|
|
)
|
|
assert len(patches) == 1
|
|
assert patches[0] == {
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/0/replicas",
|
|
"value": 2, # 1 + 1
|
|
}
|
|
|
|
def test_terminate_node(self):
|
|
self.provider.terminate(
|
|
ids=["raycluster-autoscaler-worker-small-group-dkz2r"], request_id="term-1"
|
|
)
|
|
patches = self.mock_client.get_patches(
|
|
f"rayclusters/{self.provider._cluster_name}"
|
|
)
|
|
assert len(patches) == 2
|
|
assert patches == [
|
|
{
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/0/replicas",
|
|
"value": 0,
|
|
},
|
|
{
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/0/scaleStrategy",
|
|
"value": {
|
|
"workersToDelete": [
|
|
"raycluster-autoscaler-worker-small-group-dkz2r"
|
|
]
|
|
},
|
|
},
|
|
]
|
|
|
|
def test_pending_deletes(self):
|
|
# Modify the cr.yaml to have a pending delete.
|
|
self.mock_client._ray_cluster["spec"]["workerGroupSpecs"][0][
|
|
"scaleStrategy"
|
|
] = {"workersToDelete": ["raycluster-autoscaler-worker-small-group-dkz2r"]}
|
|
self.mock_client._ray_cluster["spec"]["workerGroupSpecs"][0]["replicas"] = 0
|
|
|
|
# Launching new nodes should fail.
|
|
self.provider.launch(shape={"small-group": 1}, request_id="launch-1")
|
|
errors = self.provider.poll_errors()
|
|
assert errors[0].node_type == "small-group"
|
|
assert errors[0].request_id == "launch-1"
|
|
assert "There are workers to be deleted" in str(errors[0]), errors[0]
|
|
|
|
# Terminating new nodes should fail.
|
|
self.provider.terminate(
|
|
ids=["raycluster-autoscaler-worker-small-group-dkz2r"], request_id="term-1"
|
|
)
|
|
errors = self.provider.poll_errors()
|
|
assert (
|
|
errors[0].cloud_instance_id
|
|
== "raycluster-autoscaler-worker-small-group-dkz2r"
|
|
)
|
|
assert errors[0].request_id == "term-1"
|
|
assert "There are workers to be deleted" in str(errors[0]), errors[0]
|
|
|
|
# Remove the pod from the pod list.
|
|
self.mock_client._pod_list["items"] = [
|
|
pod
|
|
for pod in self.mock_client._pod_list["items"]
|
|
if pod["metadata"]["name"]
|
|
!= "raycluster-autoscaler-worker-small-group-dkz2r"
|
|
]
|
|
|
|
# Launch OK now, and we should also clears the pending delete.
|
|
self.provider.launch(shape={"small-group": 1}, request_id="launch-2")
|
|
errors = self.provider.poll_errors()
|
|
assert len(errors) == 0
|
|
patches = self.mock_client.get_patches(
|
|
f"rayclusters/{self.provider._cluster_name}"
|
|
)
|
|
assert len(patches) == 2
|
|
assert patches == [
|
|
{
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/0/replicas",
|
|
"value": 1,
|
|
},
|
|
{
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/0/scaleStrategy",
|
|
"value": {"workersToDelete": []},
|
|
},
|
|
]
|
|
|
|
def test_increase_min_replicas_to_scale_up(self):
|
|
# Simulate the case where users manually increase the `minReplicas` field
|
|
# from 0 to $num_pods. KubeRay will create $num_pods worker Pods to meet the new
|
|
# `minReplicas`, even though the `replicas` field is still 0.
|
|
small_group = "small-group"
|
|
num_pods = 0
|
|
assert (
|
|
self.mock_client._ray_cluster["spec"]["workerGroupSpecs"][0]["groupName"]
|
|
== small_group
|
|
)
|
|
for pod in self.mock_client._pod_list["items"]:
|
|
if pod["metadata"]["labels"]["ray.io/group"] == small_group:
|
|
num_pods += 1
|
|
assert num_pods > 0
|
|
self.mock_client._ray_cluster["spec"]["workerGroupSpecs"][0]["replicas"] = 0
|
|
self.mock_client._ray_cluster["spec"]["workerGroupSpecs"][0][
|
|
"minReplicas"
|
|
] = num_pods
|
|
|
|
# Launching a new node and `replicas` should be
|
|
# `max(replicas, minReplicas) + 1`.
|
|
self.provider.launch(shape={small_group: 1}, request_id="launch-1")
|
|
patches = self.mock_client.get_patches(
|
|
f"rayclusters/{self.provider._cluster_name}"
|
|
)
|
|
assert len(patches) == 1
|
|
assert patches[0] == {
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/0/replicas",
|
|
"value": num_pods + 1,
|
|
}
|
|
|
|
def test_inconsistent_pods_raycr_scale_up(self):
|
|
"""
|
|
Test the case where the cluster state has not yet reached the desired state.
|
|
Specifically, the replicas field in the RayCluster CR does not match the actual
|
|
number of Pods.
|
|
"""
|
|
# Check the assumptions of the test
|
|
small_group = "small-group"
|
|
num_pods = 0
|
|
for pod in self.mock_client._pod_list["items"]:
|
|
if pod["metadata"]["labels"]["ray.io/group"] == small_group:
|
|
num_pods += 1
|
|
|
|
assert (
|
|
self.mock_client._ray_cluster["spec"]["workerGroupSpecs"][0]["groupName"]
|
|
== small_group
|
|
)
|
|
desired_replicas = num_pods + 1
|
|
self.mock_client._ray_cluster["spec"]["workerGroupSpecs"][0][
|
|
"replicas"
|
|
] = desired_replicas
|
|
|
|
# Launch a new node. The replicas field should be incremented by 1, even though
|
|
# the cluster state has not yet reached the goal state.
|
|
launch_request = {"small-group": 1}
|
|
self.provider.launch(shape=launch_request, request_id="launch-1")
|
|
|
|
patches = self.mock_client.get_patches(
|
|
f"rayclusters/{self.provider._cluster_name}"
|
|
)
|
|
assert len(patches) == 1
|
|
assert patches[0] == {
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/0/replicas",
|
|
"value": desired_replicas + 1,
|
|
}
|
|
|
|
def test_inconsistent_pods_raycr_scale_down(self):
|
|
"""
|
|
Test the case where the cluster state has not yet reached the desired state.
|
|
Specifically, the replicas field in the RayCluster CR does not match the actual
|
|
number of Pods.
|
|
"""
|
|
# Check the assumptions of the test
|
|
small_group = "small-group"
|
|
num_pods = 0
|
|
pod_to_delete = None
|
|
for pod in self.mock_client._pod_list["items"]:
|
|
if pod["metadata"]["labels"]["ray.io/group"] == small_group:
|
|
num_pods += 1
|
|
pod_to_delete = pod["metadata"]["name"]
|
|
assert pod_to_delete is not None
|
|
|
|
assert (
|
|
self.mock_client._ray_cluster["spec"]["workerGroupSpecs"][0]["groupName"]
|
|
== small_group
|
|
)
|
|
desired_replicas = num_pods + 1
|
|
self.mock_client._ray_cluster["spec"]["workerGroupSpecs"][0][
|
|
"replicas"
|
|
] = desired_replicas
|
|
|
|
# Terminate a node. The replicas field should be decremented by 1, even though
|
|
# the cluster state has not yet reached the goal state.
|
|
self.provider.terminate(ids=[pod_to_delete], request_id="term-1")
|
|
patches = self.mock_client.get_patches(
|
|
f"rayclusters/{self.provider._cluster_name}"
|
|
)
|
|
assert len(patches) == 2
|
|
assert patches == [
|
|
{
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/0/replicas",
|
|
"value": desired_replicas - 1,
|
|
},
|
|
{
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/0/scaleStrategy",
|
|
"value": {
|
|
"workersToDelete": [
|
|
pod_to_delete,
|
|
]
|
|
},
|
|
},
|
|
]
|
|
|
|
def test_decrease_cr_replicas_below_observed_then_scale_down(self):
|
|
"""
|
|
If a user/operator decreases the CR's replicas below the observed number of
|
|
Pods without specifying workersToDelete, scaling down should base the
|
|
new desired on observed (floor), decrement by one, and add the pod to
|
|
workersToDelete.
|
|
"""
|
|
# Prepare a RayCluster CR with replicas set to 0 for the small-group
|
|
# while the pod list contains multiple small-group pods.
|
|
raycluster_cr = get_basic_ray_cr()
|
|
mock_client = MockKubernetesHttpApiClient(
|
|
_get_test_yaml("podlist2.yaml"), raycluster_cr
|
|
)
|
|
|
|
small_group = "small-group"
|
|
pod_names = []
|
|
for pod in mock_client._pod_list["items"]:
|
|
if pod["metadata"]["labels"]["ray.io/group"] == small_group:
|
|
pod_names.append(pod["metadata"]["name"])
|
|
assert len(pod_names) >= 2
|
|
|
|
# Decrease CR replicas below observed without workersToDelete.
|
|
assert raycluster_cr["spec"]["workerGroupSpecs"][0]["groupName"] == small_group
|
|
raycluster_cr["spec"]["workerGroupSpecs"][0]["replicas"] = 0
|
|
|
|
provider = KubeRayProvider(
|
|
cluster_name="test",
|
|
provider_config={
|
|
"namespace": "default",
|
|
"head_node_type": "headgroup",
|
|
},
|
|
gcs_client=MagicMock(),
|
|
k8s_api_client=mock_client,
|
|
)
|
|
|
|
# Terminate a single observed pod.
|
|
pod_to_delete = pod_names[0]
|
|
provider.terminate(ids=[pod_to_delete], request_id="term-decrease")
|
|
|
|
# Expected: replicas becomes observed-1; workersToDelete contains the pod.
|
|
patches = mock_client.get_patches(f"rayclusters/{provider._cluster_name}")
|
|
assert len(patches) == 2
|
|
assert patches == [
|
|
{
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/0/replicas",
|
|
"value": len(pod_names) - 1,
|
|
},
|
|
{
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/0/scaleStrategy",
|
|
"value": {
|
|
"workersToDelete": [
|
|
pod_to_delete,
|
|
]
|
|
},
|
|
},
|
|
]
|
|
|
|
def test_scale_down_multiple_pods_of_node_type(self):
|
|
"""
|
|
Test the case where multiple pods of the same node type are scaled
|
|
down on one autoscaler iteration. This test verifies that the provider
|
|
properly handles multiple pod deletions and counting workers_to_delete.
|
|
"""
|
|
# Setup provider with multiple worker pods in podlist. We use podlist2
|
|
# here because podlist1 only contains one running worker.
|
|
raycluster_cr = get_basic_ray_cr()
|
|
raycluster_cr["spec"]["workerGroupSpecs"][0]["replicas"] = 2
|
|
mock_client = MockKubernetesHttpApiClient(
|
|
_get_test_yaml("podlist2.yaml"), raycluster_cr
|
|
)
|
|
provider = KubeRayProvider(
|
|
cluster_name="test",
|
|
provider_config={
|
|
"namespace": "default",
|
|
"head_node_type": "headgroup",
|
|
},
|
|
gcs_client=MagicMock(),
|
|
k8s_api_client=mock_client,
|
|
)
|
|
|
|
# Identify all pods in the target group
|
|
small_group = "small-group"
|
|
pod_names = []
|
|
for pod in mock_client._pod_list["items"]:
|
|
if pod["metadata"]["labels"]["ray.io/group"] == small_group:
|
|
pod_names.append(pod["metadata"]["name"])
|
|
|
|
# Terminate all pods in the group
|
|
provider._sync_with_api_server()
|
|
cur_instance_ids = set(provider.instances.keys())
|
|
pods_to_terminate = [name for name in pod_names if name in cur_instance_ids]
|
|
|
|
assert (
|
|
len(pods_to_terminate) > 1
|
|
), "Expected multiple pods to terminate in the target group."
|
|
|
|
provider.terminate(ids=pods_to_terminate, request_id="term-2")
|
|
|
|
# Check the patches applied to the RayCluster resource
|
|
patches = mock_client.get_patches(f"rayclusters/{provider._cluster_name}")
|
|
|
|
assert len(patches) == 2
|
|
assert patches == [
|
|
{
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/0/replicas",
|
|
"value": 0,
|
|
},
|
|
{
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/0/scaleStrategy",
|
|
"value": {
|
|
"workersToDelete": pods_to_terminate,
|
|
},
|
|
},
|
|
]
|
|
|
|
def test_worker_to_delete_info(self):
|
|
"""
|
|
Validate _get_workers_delete_info correctly returns the worker groups with pending
|
|
deletes, worker groups with finished deletes, and the set of workers to delete.
|
|
"""
|
|
# Create a RayCluster CR and set replicas to 0 to simulate the case where the autoscaler
|
|
# patches the RayCluster with `replicas: 0`, but alive Pods still exist in workersToDelete.
|
|
raycluster_cr = get_basic_ray_cr()
|
|
raycluster_cr["spec"]["workerGroupSpecs"][0]["replicas"] = 0
|
|
mock_client = MockKubernetesHttpApiClient(
|
|
_get_test_yaml("podlist2.yaml"), raycluster_cr
|
|
)
|
|
|
|
# Add some workers to workersToDelete.
|
|
small_group = "small-group"
|
|
pod_names = []
|
|
for pod in mock_client._pod_list["items"]:
|
|
if pod["metadata"]["labels"]["ray.io/group"] == small_group:
|
|
pod_names.append(pod["metadata"]["name"])
|
|
raycluster_cr["spec"]["workerGroupSpecs"][0]["scaleStrategy"] = {
|
|
"workersToDelete": pod_names,
|
|
}
|
|
|
|
(
|
|
pending_deletes,
|
|
finished_deletes,
|
|
workers_to_delete,
|
|
) = KubeRayProvider._get_workers_delete_info(raycluster_cr, set(pod_names))
|
|
|
|
# Validate _get_workers_delete_info populates sets as expected.
|
|
assert pending_deletes == {"small-group"}
|
|
assert finished_deletes == set()
|
|
assert workers_to_delete == {pod_names[0], pod_names[1]}
|
|
|
|
def test_set_no_driver_annotation_adds_when_absent(self):
|
|
self.provider._set_no_driver_annotation()
|
|
path = f"rayclusters/{self.provider._cluster_name}"
|
|
patch = self.mock_client.get_patches(path)
|
|
assert patch == {
|
|
"metadata": {"annotations": {NO_DRIVER_TTL_EXPIRED_ANNOTATION: "true"}}
|
|
}
|
|
|
|
def test_set_no_driver_annotation_idempotent(self):
|
|
self.provider._ray_cluster.setdefault("metadata", {}).setdefault(
|
|
"annotations", {}
|
|
)[NO_DRIVER_TTL_EXPIRED_ANNOTATION] = "true"
|
|
|
|
self.provider._set_no_driver_annotation()
|
|
path = f"rayclusters/{self.provider._cluster_name}"
|
|
assert path not in self.mock_client._patches
|
|
|
|
def test_set_no_driver_annotation_swallows_patch_failure(self):
|
|
path = f"rayclusters/{self.provider._cluster_name}"
|
|
|
|
def failing_patch(*args, **kwargs):
|
|
raise RuntimeError("k8s unreachable")
|
|
|
|
self.mock_client.patch = failing_patch
|
|
# Should not raise.
|
|
self.provider._set_no_driver_annotation()
|
|
assert path not in self.mock_client._patches
|
|
|
|
# --- No-driver termination predicate + dispatch ---
|
|
|
|
def _make_gcs(self, *jobs):
|
|
class _Config:
|
|
def __init__(self, ray_namespace):
|
|
self.ray_namespace = ray_namespace
|
|
|
|
class _Job:
|
|
def __init__(self, dead, ray_namespace, end_time=0):
|
|
self.is_dead = dead
|
|
self.end_time = end_time
|
|
self.config = _Config(ray_namespace)
|
|
|
|
class _Gcs:
|
|
def get_all_job_info(self, **_):
|
|
return {i: _Job(*job) for i, job in enumerate(jobs)}
|
|
|
|
return _Gcs()
|
|
|
|
def test_driver_status_filters_internal(self):
|
|
gcs = self._make_gcs(
|
|
(False, "_ray_internal_dashboard"),
|
|
(False, "_ray_internal_something"),
|
|
)
|
|
self.provider._gcs_client = gcs
|
|
assert self.provider._driver_status()[0] is False
|
|
|
|
def test_driver_status_counts_user_driver(self):
|
|
gcs = self._make_gcs(
|
|
(False, "_ray_internal_dashboard"),
|
|
(False, "default"),
|
|
)
|
|
self.provider._gcs_client = gcs
|
|
assert self.provider._driver_status()[0] is True
|
|
|
|
def test_driver_status_ignores_dead(self):
|
|
gcs = self._make_gcs((True, "default", 42))
|
|
self.provider._gcs_client = gcs
|
|
assert self.provider._driver_status() == (False, 42)
|
|
|
|
def test_driver_status_fail_closed(self):
|
|
class _FailingGcs:
|
|
def get_all_job_info(self, **_):
|
|
raise RuntimeError("gcs unreachable")
|
|
|
|
self.provider._gcs_client = _FailingGcs()
|
|
assert self.provider._driver_status()[0] is True
|
|
|
|
def test_evaluate_no_driver_termination_disabled_when_timeout_none(self):
|
|
path = f"rayclusters/{self.provider._cluster_name}"
|
|
self.provider._gcs_client = self._make_gcs() # no drivers
|
|
self.provider._no_driver_timeout_seconds = None
|
|
self.provider._evaluate_no_driver_termination()
|
|
assert path not in self.mock_client._patches
|
|
assert self.provider._no_driver_observed_since is None
|
|
|
|
def test_evaluate_no_driver_termination_waits_for_timeout(self):
|
|
self.provider._gcs_client = self._make_gcs() # no drivers
|
|
self.provider._no_driver_timeout_seconds = 100.0
|
|
|
|
def evaluate_at(t):
|
|
with mock.patch("time.monotonic", return_value=t):
|
|
self.provider._evaluate_no_driver_termination()
|
|
|
|
path = f"rayclusters/{self.provider._cluster_name}"
|
|
evaluate_at(0.0)
|
|
assert path not in self.mock_client._patches # anchored, not yet
|
|
evaluate_at(50.0)
|
|
assert path not in self.mock_client._patches # still below timeout
|
|
evaluate_at(100.0)
|
|
assert self.mock_client._patches.get(path) == {
|
|
"metadata": {"annotations": {NO_DRIVER_TTL_EXPIRED_ANNOTATION: "true"}}
|
|
}
|
|
|
|
def test_evaluate_no_driver_termination_resets_when_driver_attaches(self):
|
|
path = f"rayclusters/{self.provider._cluster_name}"
|
|
self.provider._no_driver_timeout_seconds = 100.0
|
|
|
|
with mock.patch("time.monotonic", return_value=0.0):
|
|
self.provider._gcs_client = self._make_gcs() # no drivers
|
|
self.provider._evaluate_no_driver_termination()
|
|
assert self.provider._no_driver_observed_since == 0.0
|
|
|
|
# Driver attaches → anchor cleared, no patch.
|
|
with mock.patch("time.monotonic", return_value=50.0):
|
|
self.provider._gcs_client = self._make_gcs((False, "default"))
|
|
self.provider._evaluate_no_driver_termination()
|
|
assert self.provider._no_driver_observed_since is None
|
|
assert path not in self.mock_client._patches
|
|
|
|
def test_evaluate_no_driver_termination_resets_on_intermittent_driver(self):
|
|
self.provider._no_driver_timeout_seconds = 100.0
|
|
|
|
# No driver: anchor at t=0.
|
|
with mock.patch("time.monotonic", return_value=0.0):
|
|
self.provider._gcs_client = self._make_gcs()
|
|
self.provider._evaluate_no_driver_termination()
|
|
assert self.provider._no_driver_observed_since == 0.0
|
|
|
|
# A short-lived driver started and finished between loops (a dead job
|
|
# with a newer end time): the timer must restart.
|
|
with mock.patch("time.monotonic", return_value=50.0):
|
|
self.provider._gcs_client = self._make_gcs((True, "default", 42))
|
|
self.provider._evaluate_no_driver_termination()
|
|
assert self.provider._no_driver_observed_since == 50.0
|
|
assert self.provider._last_seen_job_end_time == 42
|
|
|
|
def test_refresh_no_driver_timeout_seconds_reads_value(self):
|
|
self.provider._ray_cluster = {
|
|
"spec": {"autoscalerOptions": {"noDriverTimeoutSeconds": 1800}}
|
|
}
|
|
self.provider._refresh_no_driver_timeout_seconds()
|
|
assert self.provider._no_driver_timeout_seconds == 1800.0
|
|
|
|
def test_refresh_no_driver_timeout_seconds_unset(self):
|
|
self.provider._ray_cluster = {"spec": {"autoscalerOptions": {}}}
|
|
self.provider._refresh_no_driver_timeout_seconds()
|
|
assert self.provider._no_driver_timeout_seconds is None
|
|
|
|
def test_refresh_no_driver_timeout_seconds_no_autoscaler_options(self):
|
|
self.provider._ray_cluster = {"spec": {}}
|
|
self.provider._refresh_no_driver_timeout_seconds()
|
|
assert self.provider._no_driver_timeout_seconds is None
|
|
|
|
def test_scale_down_with_multi_host_group(self):
|
|
"""
|
|
Test the case where a worker group has numOfHosts > 1.
|
|
This ensures that the KubeRay provider accounts for multi-host replicas
|
|
during scale down and properly updates the workersToDelete field.
|
|
"""
|
|
# Setup mock RayCluster CR with numOfHosts: 2 and replicas: 1,
|
|
# resulting in 2 workers total.
|
|
raycluster_cr = get_basic_ray_cr()
|
|
raycluster_cr["spec"]["workerGroupSpecs"][0]["replicas"] = 2
|
|
mock_client = MockKubernetesHttpApiClient(
|
|
_get_test_yaml("podlist2.yaml"), raycluster_cr
|
|
)
|
|
provider = KubeRayProvider(
|
|
cluster_name="test",
|
|
provider_config={
|
|
"namespace": "default",
|
|
"head_node_type": "headgroup",
|
|
},
|
|
gcs_client=MagicMock(),
|
|
k8s_api_client=mock_client,
|
|
)
|
|
|
|
# Identify all pods in the multi-host group
|
|
pod_names = []
|
|
for pod in mock_client._pod_list["items"]:
|
|
if pod["metadata"]["labels"]["ray.io/group"] == "tpu-group":
|
|
pod_names.append(pod["metadata"]["name"])
|
|
|
|
# Expect 2 pods since replicas=1 and numOfHosts=2
|
|
assert len(pod_names) == 2, "Expected 2 pods in the multi-host group."
|
|
|
|
# Sync provider state and mark all pods for deletion
|
|
provider._sync_with_api_server()
|
|
cur_instance_ids = set(provider.instances.keys())
|
|
pods_to_terminate = [name for name in pod_names if name in cur_instance_ids]
|
|
|
|
assert (
|
|
len(pods_to_terminate) == 2
|
|
), "Expected all multi-host pods to be tracked by the provider."
|
|
|
|
# Terminate all pods in the group
|
|
provider.terminate(ids=pods_to_terminate, request_id="term-multi")
|
|
|
|
# Check that scale request successfully created
|
|
patches = mock_client.get_patches(f"rayclusters/{provider._cluster_name}")
|
|
|
|
assert len(patches) == 2
|
|
assert patches == [
|
|
{
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/2/replicas",
|
|
"value": 0,
|
|
},
|
|
{
|
|
"op": "replace",
|
|
"path": "/spec/workerGroupSpecs/2/scaleStrategy",
|
|
"value": {
|
|
"workersToDelete": pods_to_terminate,
|
|
},
|
|
},
|
|
]
|
|
|
|
|
|
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__]))
|