wehub-resource-sync
666 lines
25 KiB
Python
666 lines
25 KiB
Python
import copy
|
|
import os
|
|
import re
|
|
import sys
|
|
import tempfile
|
|
from threading import RLock
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from ray.autoscaler._private.vsphere.cluster_operator_client import (
|
|
ClusterOperatorClient,
|
|
VMNodeStatus,
|
|
)
|
|
from ray.autoscaler.tags import (
|
|
NODE_KIND_HEAD,
|
|
NODE_KIND_WORKER,
|
|
STATUS_SETTING_UP,
|
|
STATUS_UNINITIALIZED,
|
|
STATUS_UP_TO_DATE,
|
|
TAG_RAY_LAUNCH_CONFIG,
|
|
TAG_RAY_NODE_KIND,
|
|
TAG_RAY_NODE_NAME,
|
|
TAG_RAY_NODE_STATUS,
|
|
TAG_RAY_USER_NODE_TYPE,
|
|
)
|
|
|
|
_CLUSTER_NAME = "ray-cluster"
|
|
_PROVIDER_CONFIG = {
|
|
"vsphere_config": {
|
|
"namespace": "deploy-ray",
|
|
"ca_cert": "",
|
|
"api_server": "10.10.10.10",
|
|
}
|
|
}
|
|
|
|
_CLUSTER_CONFIG = {
|
|
"cluster_name": "ray-cluster",
|
|
"max_workers": 5,
|
|
"upscaling_speed": 1.0,
|
|
"docker": {
|
|
"image": "project-taiga-docker-local.artifactory.eng.vmware.com/\
|
|
development/ray:milestone_2",
|
|
"container_name": "ray_container",
|
|
"pull_before_run": True,
|
|
"run_options": [],
|
|
},
|
|
"idle_timeout_minutes": 5,
|
|
"provider": {
|
|
"type": "vsphere",
|
|
"vsphere_config": {
|
|
"namespace": "deploy-ray",
|
|
"ca_cert": "",
|
|
"api_server": "10.78.116.4",
|
|
"vm_image": "vmi-86621c422637d11ef",
|
|
"storage_class": "wcpe2e-nfs-profile",
|
|
},
|
|
},
|
|
"auth": {
|
|
"ssh_user": "ray",
|
|
"ssh_private_key": os.path.join(tempfile.mkdtemp(), "id_rsa_ray.pem"),
|
|
},
|
|
"available_node_types": {
|
|
"ray.head.default": {
|
|
"resources": {},
|
|
"node_config": {"vm_class": "best-effort-xlarge"},
|
|
},
|
|
"worker": {
|
|
"min_workers": 1,
|
|
"max_workers": 3,
|
|
"resources": {},
|
|
"node_config": {"vm_class": "best-effort-xlarge"},
|
|
},
|
|
},
|
|
"head_node_type": "ray.head.default",
|
|
"file_mounts": {},
|
|
"cluster_synced_files": [],
|
|
"file_mounts_sync_continuously": False,
|
|
"rsync_exclude": [],
|
|
"rsync_filter": [],
|
|
"initialization_commands": [],
|
|
"setup_commands": [],
|
|
"head_setup_commands": {},
|
|
"worker_setup_commands": [],
|
|
"head_start_ray_commands": [],
|
|
"worker_start_ray_commands": [],
|
|
}
|
|
|
|
|
|
_CLUSTER_RESPONSE = {
|
|
"metadata": {
|
|
"creationTimestamp": "2024-08-01T07:20:42Z",
|
|
"deletionGracePeriodSeconds": 0,
|
|
"deletionTimestamp": "2024-08-01T07:30:14Z",
|
|
"finalizers": ["vmraycluster.vmray.broadcom.com"],
|
|
"generation": 3,
|
|
"labels": {"vmray.kubernetes.io/head-nounce": "ly2uv"},
|
|
},
|
|
"spec": {
|
|
"api_server": {"location": "10.185.109.130"},
|
|
"autoscaler_desired_workers": {
|
|
"ray-cluster-w-8bx6j4n1": "ray.worker.default",
|
|
"ray-cluster-w-nb9u9xf3": "ray.worker.default",
|
|
},
|
|
"common_node_config": {
|
|
"available_node_types": {
|
|
"ray.head.default": {
|
|
"max_workers": 0,
|
|
"min_workers": 0,
|
|
"resources": {},
|
|
"vm_class": "best-effort-xlarge",
|
|
},
|
|
"ray.worker.default": {
|
|
"max_workers": 5,
|
|
"min_workers": 2,
|
|
"resources": {"cpu": 4, "memory": 4294967296},
|
|
"vm_class": "best-effort-xlarge",
|
|
},
|
|
},
|
|
"max_workers": 5,
|
|
"storage_class": "vsan-default-storage-policy",
|
|
"vm_image": "vmi-76e37912125b9ad17",
|
|
"vm_password_salt_hash": "$6$test1234$9/BUZHNkvq.c1miDDMG5cHLmM4V7gbYdGuF0/\
|
|
/3gSIh//DOyi7ypPCs6EAA9b8/tidHottL6UG0tG/RqTgAAi/",
|
|
"vm_user": "ray-vm",
|
|
},
|
|
"enable_tls": True,
|
|
"head_node": {"port": 6254, "node_type": "ray.head.default"},
|
|
"ray_docker_image": "project-taiga-docker-local.artifactory.eng.vmware.com/\
|
|
development/ray:milestone_2.0",
|
|
},
|
|
"status": {
|
|
"cluster_state": "healthy",
|
|
"current_workers": {
|
|
"ray-cluster-w-8bx6j4n1": {"ip": "10.244.1.131", "vm_status": "setting-up"},
|
|
"ray-cluster-w-nb9u9xf3": {"ip": "10.244.1.132", "vm_status": "setting-up"},
|
|
},
|
|
"head_node_status": {
|
|
"ip": "10.244.1.130",
|
|
"ray_status": "running",
|
|
"vm_status": "running",
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def create_random_pvt_key():
|
|
from cryptography.hazmat.primitives import serialization
|
|
from cryptography.hazmat.primitives.asymmetric import rsa
|
|
|
|
private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
|
pem_private_key = private_key.private_bytes(
|
|
encoding=serialization.Encoding.PEM,
|
|
format=serialization.PrivateFormat.PKCS8,
|
|
encryption_algorithm=serialization.NoEncryption(),
|
|
)
|
|
|
|
delete_random_pvt_key()
|
|
file = open(_CLUSTER_CONFIG["auth"]["ssh_private_key"], "w")
|
|
file.write(pem_private_key.decode())
|
|
file.close()
|
|
|
|
|
|
def delete_random_pvt_key():
|
|
filename = _CLUSTER_CONFIG["auth"]["ssh_private_key"]
|
|
if os.path.exists(filename):
|
|
os.remove(filename)
|
|
|
|
|
|
def mock_cluster_operator():
|
|
def __init__(self, cluster_name: str, provider_config: dict, cluster_config: dict):
|
|
self.cluster_name = cluster_name
|
|
self.vmraycluster_nounce = None
|
|
self.max_worker_nodes = None
|
|
self.vsphere_config = provider_config["vsphere_config"]
|
|
self.namespace = self.vsphere_config["namespace"]
|
|
os.environ["SERVICE_ACCOUNT_TOKEN"] = "test_account"
|
|
self.k8s_api_client = MagicMock()
|
|
if cluster_config:
|
|
self.max_worker_nodes = cluster_config["max_workers"]
|
|
self.head_setup_commands = cluster_config["head_setup_commands"]
|
|
self.available_node_types = cluster_config["available_node_types"]
|
|
self.head_node_type = cluster_config["head_node_type"]
|
|
|
|
# docker configurations.
|
|
self.provider_auth = cluster_config["auth"]
|
|
self.docker = cluster_config["docker"]
|
|
self.docker_config = None
|
|
|
|
with patch.object(ClusterOperatorClient, "__init__", __init__):
|
|
operator = ClusterOperatorClient(
|
|
_CLUSTER_NAME, _PROVIDER_CONFIG, _CLUSTER_CONFIG
|
|
)
|
|
return copy.deepcopy(operator)
|
|
|
|
|
|
def test_list_vms():
|
|
"""Should return nodes as per tags provided in the tag_filters."""
|
|
operator = mock_cluster_operator()
|
|
operator.lock = RLock()
|
|
head_node_tags = {TAG_RAY_NODE_KIND: NODE_KIND_HEAD}
|
|
head_node_tags[TAG_RAY_LAUNCH_CONFIG] = "launch_hash"
|
|
head_node_tags[TAG_RAY_NODE_NAME] = "ray-{}-head".format(operator.cluster_name)
|
|
head_node_tags[TAG_RAY_NODE_STATUS] = STATUS_UNINITIALIZED
|
|
# Mock cluster response
|
|
operator._get_cluster_response = MagicMock(return_value=_CLUSTER_RESPONSE)
|
|
head_node = f"{operator.cluster_name}-h-1234"
|
|
operator._get_head_name = MagicMock(return_value=head_node)
|
|
# make sure function provides head node
|
|
nodes, tag_cache = operator.list_vms(head_node_tags)
|
|
# make sure max number of worker nodes are set up
|
|
assert operator.max_worker_nodes == 5
|
|
assert len(nodes) == 1
|
|
assert nodes[0] == head_node
|
|
assert head_node in tag_cache.keys()
|
|
assert tag_cache[head_node][TAG_RAY_NODE_NAME] == head_node
|
|
assert tag_cache[head_node][TAG_RAY_NODE_STATUS] == STATUS_UP_TO_DATE
|
|
assert tag_cache[head_node][TAG_RAY_NODE_KIND] == NODE_KIND_HEAD
|
|
assert tag_cache[head_node][TAG_RAY_USER_NODE_TYPE] == "ray.head.default"
|
|
# In case of wrong node kind no node should be returned
|
|
node_tags = head_node_tags.copy()
|
|
node_tags[TAG_RAY_NODE_KIND] = "wrong_node"
|
|
nodes, tag_cache = operator.list_vms(node_tags)
|
|
assert len(nodes) == 0
|
|
assert len(tag_cache) == 0
|
|
|
|
# Make sure status updated properly
|
|
cluster_response = _CLUSTER_RESPONSE.copy()
|
|
cluster_response["status"]["head_node_status"][
|
|
"vm_status"
|
|
] = VMNodeStatus.INITIALIZED.value
|
|
# Mock cluster response
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
# make sure the function provides head node only
|
|
nodes, tag_cache = operator.list_vms(head_node_tags)
|
|
assert nodes[0] == head_node
|
|
assert head_node in tag_cache.keys()
|
|
assert tag_cache[head_node][TAG_RAY_NODE_NAME] == head_node
|
|
assert tag_cache[head_node][TAG_RAY_NODE_STATUS] == STATUS_SETTING_UP
|
|
|
|
cluster_response["status"]["head_node_status"][
|
|
"vm_status"
|
|
] = VMNodeStatus.FAIL.value
|
|
# Mock cluster response
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
# make sure the function provides head node only
|
|
nodes, tag_cache = operator.list_vms(head_node_tags)
|
|
assert nodes[0] == head_node
|
|
assert head_node in tag_cache.keys()
|
|
assert tag_cache[head_node][TAG_RAY_NODE_NAME] == head_node
|
|
assert tag_cache[head_node][TAG_RAY_NODE_STATUS] == STATUS_UNINITIALIZED
|
|
|
|
# make sure the function returns only worker nodes
|
|
worker_node_tags = {TAG_RAY_NODE_KIND: NODE_KIND_WORKER}
|
|
cluster_response = _CLUSTER_RESPONSE.copy()
|
|
# Mock cluster response
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
nodes, tag_cache = operator.list_vms(worker_node_tags)
|
|
assert len(nodes) == 2
|
|
for node in nodes:
|
|
assert tag_cache[node][TAG_RAY_NODE_NAME] == node
|
|
assert tag_cache[node][TAG_RAY_NODE_STATUS] == STATUS_UNINITIALIZED
|
|
assert tag_cache[node][TAG_RAY_NODE_KIND] == NODE_KIND_WORKER
|
|
|
|
# In case of no tag_filters provided return head node as well as worker nodes
|
|
tag_filters = {}
|
|
cluster_response = _CLUSTER_RESPONSE.copy()
|
|
cluster_response["status"]["head_node_status"][
|
|
"vm_status"
|
|
] = VMNodeStatus.RUNNING.value
|
|
# Mock cluster response
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
# make sure the function provides head node as well all worker nodes.
|
|
nodes, tag_cache = operator.list_vms(tag_filters)
|
|
assert len(nodes) == 3
|
|
assert head_node in tag_cache.keys()
|
|
assert tag_cache[head_node][TAG_RAY_NODE_NAME] == head_node
|
|
assert tag_cache[head_node][TAG_RAY_NODE_STATUS] == STATUS_UP_TO_DATE
|
|
for node in nodes:
|
|
if node == head_node:
|
|
continue
|
|
assert tag_cache[node][TAG_RAY_NODE_NAME] == node
|
|
assert tag_cache[node][TAG_RAY_NODE_STATUS] == STATUS_UNINITIALIZED
|
|
assert tag_cache[node][TAG_RAY_NODE_KIND] == NODE_KIND_WORKER
|
|
# make sure the function returns nodes which are not yet initialised but are in
|
|
# desired workers list
|
|
cluster_response = _CLUSTER_RESPONSE.copy()
|
|
del cluster_response["status"]["current_workers"]["ray-cluster-w-nb9u9xf3"]
|
|
# Mock cluster response
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
nodes, tag_cache = operator.list_vms(tag_filters)
|
|
desired_worker = "ray-cluster-w-nb9u9xf3"
|
|
assert len(nodes) == 3
|
|
assert tag_cache[desired_worker][TAG_RAY_NODE_NAME] == desired_worker
|
|
assert tag_cache[desired_worker][TAG_RAY_NODE_STATUS] == STATUS_UNINITIALIZED
|
|
assert tag_cache[desired_worker][TAG_RAY_NODE_KIND] == NODE_KIND_WORKER
|
|
|
|
|
|
def test_is_vm_power_on():
|
|
"""Should return True only if VM is in Running state"""
|
|
operator = mock_cluster_operator()
|
|
operator.lock = RLock()
|
|
operator._get_node = MagicMock(
|
|
return_value={"vm_status": VMNodeStatus.RUNNING.value}
|
|
)
|
|
assert operator.is_vm_power_on("test_vm") is True
|
|
|
|
operator._get_node = MagicMock(
|
|
return_value={"vm_status": VMNodeStatus.INITIALIZED.value}
|
|
)
|
|
assert operator.is_vm_power_on("test_vm") is False
|
|
|
|
|
|
def test_is_vm_creating():
|
|
"""Should return True only if VM is in INITIALIZED state"""
|
|
operator = mock_cluster_operator()
|
|
operator.lock = RLock()
|
|
operator._get_node = MagicMock(
|
|
return_value={"vm_status": VMNodeStatus.RUNNING.value}
|
|
)
|
|
assert operator.is_vm_creating("test_vm") is False
|
|
operator._get_node = MagicMock(
|
|
return_value={"vm_status": VMNodeStatus.INITIALIZED.value}
|
|
)
|
|
assert operator.is_vm_creating("test_vm") is True
|
|
|
|
|
|
def test_get_vm_external_ip():
|
|
"""Should return IP if it is a valid one otherwise return None"""
|
|
operator = mock_cluster_operator()
|
|
# Check for worker node
|
|
operator._get_head_name = MagicMock(return_value=f"{operator.cluster_name}-h-1234")
|
|
operator.lock = RLock()
|
|
# Check IP is valid one
|
|
operator._get_node = MagicMock(
|
|
return_value={"vm_status": VMNodeStatus.RUNNING.value, "ip": "10.10.10.10"}
|
|
)
|
|
ip = operator.get_vm_external_ip("test_vm")
|
|
assert ip == "10.10.10.10"
|
|
# Check for invalid IP
|
|
operator._get_node = MagicMock(
|
|
return_value={"vm_status": VMNodeStatus.RUNNING.value, "ip": ""}
|
|
)
|
|
ip = operator.get_vm_external_ip("test_vm")
|
|
assert ip is None
|
|
# Check for invalid IP
|
|
operator._get_node = MagicMock(
|
|
return_value={"vm_status": VMNodeStatus.FAIL.value, "ip": ""}
|
|
)
|
|
ip = operator.get_vm_external_ip("test_vm")
|
|
assert ip is None
|
|
|
|
# Test head node
|
|
operator._get_vm_service_ingress = MagicMock(return_value=[{"ip": "10.20.10.10"}])
|
|
ip = operator.get_vm_external_ip(f"{operator.cluster_name}-h-1234")
|
|
assert ip == "10.20.10.10"
|
|
|
|
operator._get_vm_service_ingress = MagicMock(return_value=[{"ip": ""}])
|
|
ip = operator.get_vm_external_ip(f"{operator.cluster_name}-h-1234")
|
|
assert ip is None
|
|
|
|
operator._get_vm_service_ingress = MagicMock(return_value=[{}])
|
|
ip = operator.get_vm_external_ip(f"{operator.cluster_name}-h-1234")
|
|
assert ip is None
|
|
|
|
|
|
def test_delete_node():
|
|
"""Should call _patch with updated desired worker nodes"""
|
|
operator = mock_cluster_operator()
|
|
operator.lock = RLock()
|
|
cluster_response = _CLUSTER_RESPONSE.copy()
|
|
# Mock cluster response
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
operator.k8s_api_client.custom_object_api.patch_namespaced_custom_object = (
|
|
MagicMock()
|
|
)
|
|
operator.delete_node("ray-cluster-w-nb9u9xf3")
|
|
custom_api = operator.k8s_api_client.custom_object_api
|
|
custom_api.patch_namespaced_custom_object.assert_called_once()
|
|
|
|
|
|
def test_create_nodes():
|
|
"""Should call _patch to create worker nodes"""
|
|
operator = mock_cluster_operator()
|
|
operator.lock = RLock()
|
|
cluster_response = _CLUSTER_RESPONSE.copy()
|
|
operator.k8s_api_client.custom_object_api.patch_namespaced_custom_object = (
|
|
MagicMock()
|
|
)
|
|
# Mock cluster response
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
# Make sure API is not getting called to create a head node as it should be
|
|
# created by the operator.
|
|
tags = {TAG_RAY_NODE_NAME: "head", TAG_RAY_USER_NODE_TYPE: "ray.default.head"}
|
|
operator._patch = MagicMock()
|
|
operator._create_secret = MagicMock()
|
|
operator._create_vmraycluster = MagicMock()
|
|
head_node_name = f"{operator.cluster_name}-h-1234"
|
|
operator._create_node_name = MagicMock(return_value=head_node_name)
|
|
create_random_pvt_key()
|
|
created_nodes = operator.create_nodes(tags, 1, {})
|
|
delete_random_pvt_key()
|
|
assert len(created_nodes) == 1
|
|
assert created_nodes[head_node_name] == head_node_name
|
|
operator._create_secret.assert_called_once()
|
|
operator._create_vmraycluster.assert_called_once()
|
|
# if nodes to be creaed is zero, _patch should not be called
|
|
operator._create_secret = MagicMock()
|
|
operator._create_vmraycluster = MagicMock()
|
|
created_nodes = operator.create_nodes(tags, 0, {})
|
|
assert len(created_nodes) == 0
|
|
operator._create_secret.assert_not_called()
|
|
operator._create_vmraycluster.assert_not_called()
|
|
# Make sure desired number of workers get created
|
|
tags = {TAG_RAY_NODE_NAME: "worker", TAG_RAY_USER_NODE_TYPE: "ray.default.worker"}
|
|
operator.max_worker_nodes = 5
|
|
operator._get_cluster_response = MagicMock(return_value={})
|
|
operator._create_node_name = MagicMock()
|
|
operator.k8s_api_client.custom_object_api.patch_namespaced_custom_object = (
|
|
MagicMock()
|
|
)
|
|
operator._create_node_name.side_effect = ["test-w-1", "test-w-2"]
|
|
created_nodes = operator.create_nodes(tags, 2, {})
|
|
custom_api = operator.k8s_api_client.custom_object_api
|
|
custom_api.patch_namespaced_custom_object.assert_called()
|
|
assert len(created_nodes) == 2
|
|
# Make sure the function creates worker nodes if no desired workers present
|
|
cluster_response = _CLUSTER_RESPONSE.copy()
|
|
cluster_response["spec"]["autoscaler_desired_workers"] = {}
|
|
operator._create_node_name = MagicMock()
|
|
operator._create_node_name.side_effect = [
|
|
"test-w-1",
|
|
"test-w-2",
|
|
"test-w-3",
|
|
"test-w-4",
|
|
"test-w-5",
|
|
]
|
|
# Mock cluster response
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
operator.k8s_api_client.custom_object_api.patch_namespaced_custom_object = (
|
|
MagicMock()
|
|
)
|
|
created_nodes = operator.create_nodes(tags, 5, {})
|
|
custom_api = operator.k8s_api_client.custom_object_api
|
|
custom_api.patch_namespaced_custom_object.assert_called()
|
|
assert len(created_nodes) == 5
|
|
# Make sure not to overprovision the worker nodes
|
|
cluster_response = _CLUSTER_RESPONSE.copy()
|
|
cluster_response["spec"]["autoscaler_desired_workers"] = {
|
|
"test-1": "ray.default.worker",
|
|
"test-2": "ray.default.worker",
|
|
"test-3": "ray.default.worker",
|
|
}
|
|
operator._create_node_name = MagicMock()
|
|
operator._create_node_name.side_effect = ["test-w-4", "test-w-5", "test-w-6"]
|
|
# Mock cluster response
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
operator.k8s_api_client.custom_object_api.patch_namespaced_custom_object = (
|
|
MagicMock()
|
|
)
|
|
created_nodes = operator.create_nodes(tags, 3, {})
|
|
custom_api = operator.k8s_api_client.custom_object_api
|
|
custom_api.patch_namespaced_custom_object.assert_not_called()
|
|
assert len(created_nodes) == 0
|
|
|
|
|
|
def test__get_node():
|
|
"""Should provide a node as per given node_id"""
|
|
operator = mock_cluster_operator()
|
|
operator.lock = RLock()
|
|
cluster_response = _CLUSTER_RESPONSE.copy()
|
|
# Mock cluster response
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
# Get head node
|
|
head_node = f"{operator.cluster_name}-head"
|
|
operator._get_head_name = MagicMock(return_value=head_node)
|
|
node = operator._get_node(head_node)
|
|
assert node["vm_status"] == VMNodeStatus.RUNNING.value
|
|
assert node["ip"] == "10.244.1.130"
|
|
# Get worker node from the current workers
|
|
cluster_response["status"]["current_workers"] = {
|
|
"ray-cluster-w-8bx6j4n1": {"vm_status": "initialized"}
|
|
}
|
|
worker_node = "ray-cluster-w-8bx6j4n1"
|
|
node = operator._get_node(worker_node)
|
|
assert node["vm_status"] == VMNodeStatus.INITIALIZED.value
|
|
# get worker from the desired workers
|
|
cluster_response["spec"]["autoscaler_desired_workers"] = {
|
|
"test-vm1": "ray.default.worker"
|
|
}
|
|
# Mock cluster response
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
worker_node = "test-vm1"
|
|
node = operator._get_node(worker_node)
|
|
assert node["vm_status"] == VMNodeStatus.INITIALIZED.value
|
|
# Should not return any node
|
|
cluster_response = _CLUSTER_RESPONSE.copy()
|
|
# Mock cluster response
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
worker_node = "test-vm4"
|
|
assert operator._get_node(worker_node) == {}
|
|
|
|
|
|
def test_safe_to_scale():
|
|
"""Should not scale until current state matches desired state"""
|
|
operator = mock_cluster_operator()
|
|
operator.lock = RLock()
|
|
operator.max_worker_nodes = 5
|
|
cluster_response = _CLUSTER_RESPONSE.copy()
|
|
# Case where not all workers are up and running
|
|
cluster_response["status"]["current_workers"] = {
|
|
"test-vm1": {"vm_status": "running"}
|
|
}
|
|
cluster_response["spec"]["autoscaler_desired_workers"] = {
|
|
"test-vm1": "ray.default.worker",
|
|
"test-vm2": "ray.default.worker",
|
|
}
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
assert operator.safe_to_scale() is False
|
|
# Case where workers are waiting to be deleted
|
|
cluster_response["status"]["current_workers"] = {
|
|
"test-vm1": {"vm_status": "running"},
|
|
"test-vm2": {"vm_status": "running"},
|
|
}
|
|
cluster_response["spec"]["autoscaler_desired_workers"] = {
|
|
"test-vm1": "ray.default.worker"
|
|
}
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
assert operator.safe_to_scale() is False
|
|
# Case where current state is a desired state
|
|
operator._patch = MagicMock()
|
|
cluster_response["spec"]["autoscaler_desired_workers"] = {
|
|
"test-vm1": "ray.default.worker",
|
|
"test-vm2": "ray.default.worker",
|
|
}
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
assert operator.safe_to_scale() is True
|
|
|
|
|
|
def test__create_node_name():
|
|
"""Should create node name for head and worker nodes"""
|
|
operator = mock_cluster_operator()
|
|
# test for head node name
|
|
pattern = re.compile("^(ray-cluster-h-[a-z0-9]{8})$")
|
|
node_name = operator._create_node_name("ray-head-node")
|
|
assert pattern.match(node_name) is not None
|
|
# test for worker node name
|
|
pattern = re.compile("^(ray-cluster-w-[a-z0-9]{8})$")
|
|
node_name = operator._create_node_name("ray-worker-node")
|
|
assert pattern.match(node_name) is not None
|
|
|
|
|
|
def test__get_head_name():
|
|
"""Should return node name for a head node"""
|
|
operator = mock_cluster_operator()
|
|
operator.vmraycluster_nounce = None
|
|
# If nounce is not set read it from
|
|
# cluster response
|
|
cluster_response = _CLUSTER_RESPONSE.copy()
|
|
nounce = cluster_response["metadata"]["labels"]["vmray.kubernetes.io/head-nounce"]
|
|
# Mock cluster response
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
head_node_name = operator._get_head_name()
|
|
assert head_node_name == f"{operator.cluster_name}-h-{nounce}"
|
|
|
|
|
|
def test__create_vmraycluster():
|
|
"""Should create VMRayCluster CRD"""
|
|
operator = mock_cluster_operator()
|
|
operator.vmraycluster_nounce = "1234"
|
|
operator.k8s_api_client.custom_object_api.create_namespaced_custom_object = (
|
|
MagicMock()
|
|
)
|
|
operator._create_vmraycluster()
|
|
custom_api = operator.k8s_api_client.custom_object_api
|
|
custom_api.create_namespaced_custom_object.assert_called_once()
|
|
|
|
|
|
def test__set_tags():
|
|
"""Should set tags for a given node"""
|
|
operator = mock_cluster_operator()
|
|
node_id = "test-1"
|
|
node_kind = "head"
|
|
node_user_type = "ray.default.head"
|
|
node_status = "running"
|
|
new_tags = operator._set_tags(
|
|
node_id,
|
|
node_kind,
|
|
node_user_type,
|
|
node_status,
|
|
{
|
|
"ray-node-status": "",
|
|
"ray-node-name": "",
|
|
"ray-node-type": "",
|
|
"ray-user-node-type": "",
|
|
"test": "test",
|
|
},
|
|
)
|
|
assert new_tags[TAG_RAY_NODE_STATUS] == STATUS_UP_TO_DATE
|
|
assert new_tags[TAG_RAY_NODE_NAME] == node_id
|
|
assert new_tags[TAG_RAY_NODE_KIND] == node_kind
|
|
assert new_tags[TAG_RAY_USER_NODE_TYPE] == node_user_type
|
|
assert new_tags["test"] == "test"
|
|
node_status = "initialized"
|
|
new_tags = operator._set_tags(
|
|
node_id,
|
|
node_kind,
|
|
node_user_type,
|
|
node_status,
|
|
{
|
|
"ray-node-status": "",
|
|
"ray-node-name": "",
|
|
"ray-node-type": "",
|
|
"ray-user-node-type": "",
|
|
"test": "test",
|
|
},
|
|
)
|
|
assert new_tags[TAG_RAY_NODE_STATUS] == STATUS_SETTING_UP
|
|
node_status = "any_other_status"
|
|
new_tags = operator._set_tags(
|
|
node_id,
|
|
node_kind,
|
|
node_user_type,
|
|
node_status,
|
|
{
|
|
"ray-node-status": "",
|
|
"ray-node-name": "",
|
|
"ray-node-type": "",
|
|
"ray-user-node-type": "",
|
|
"test": "test",
|
|
},
|
|
)
|
|
assert new_tags[TAG_RAY_NODE_STATUS] == STATUS_UNINITIALIZED
|
|
|
|
|
|
def test__set_max_worker_nodes():
|
|
operator = mock_cluster_operator()
|
|
cluster_response = _CLUSTER_RESPONSE.copy()
|
|
operator._get_cluster_response = MagicMock(return_value=cluster_response)
|
|
operator.max_worker_nodes = None
|
|
operator._set_max_worker_nodes()
|
|
assert operator.max_worker_nodes == 5
|
|
|
|
|
|
def test_get_vm_service_ingress():
|
|
operator = mock_cluster_operator()
|
|
vm_service_response = {
|
|
"status": {"loadBalancer": {"ingress": [{"ip": "10.10.10.10"}]}}
|
|
}
|
|
operator._get_vm_service = MagicMock(return_value=vm_service_response)
|
|
ingress = operator._get_vm_service_ingress()
|
|
assert len(ingress) == 1
|
|
assert ingress[0]["ip"] == "10.10.10.10"
|
|
# Check if status is not available
|
|
vm_service_response = {}
|
|
operator._get_vm_service = MagicMock(return_value=vm_service_response)
|
|
assert operator._get_vm_service_ingress() == []
|
|
# Check if ingress in unavailable
|
|
vm_service_response = {"status": {"loadBalancer": {}}}
|
|
operator._get_vm_service = MagicMock(return_value=vm_service_response)
|
|
assert operator._get_vm_service_ingress() == []
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|