wehub-resource-sync
445 lines
14 KiB
Python
445 lines
14 KiB
Python
import logging
|
|
import sys
|
|
from threading import RLock
|
|
from typing import Dict
|
|
from unittest.mock import MagicMock, call, patch
|
|
|
|
import pytest
|
|
|
|
from ray.autoscaler._private.command_runner import DockerCommandRunner, SSHCommandRunner
|
|
from ray.autoscaler._private.gcp.config import (
|
|
_get_num_tpu_chips,
|
|
_has_tpus_in_node_configs,
|
|
_is_single_host_tpu,
|
|
get_node_type,
|
|
tpu_accelerator_config_to_type,
|
|
)
|
|
from ray.autoscaler._private.gcp.node import (
|
|
GCPCompute,
|
|
GCPNode,
|
|
GCPNodeType,
|
|
GCPResource,
|
|
)
|
|
from ray.autoscaler._private.gcp.node_provider import GCPNodeProvider
|
|
from ray.autoscaler._private.gcp.tpu_command_runner import (
|
|
TPUCommandRunner,
|
|
TPUVMDockerCommandRunner,
|
|
TPUVMSSHCommandRunner,
|
|
)
|
|
from ray.tests.test_autoscaler import MockProcessRunner
|
|
|
|
_PROJECT_NAME = "project-one"
|
|
_AZ = "us-west1-b"
|
|
|
|
auth_config = {
|
|
"ssh_user": "ray",
|
|
"ssh_private_key": "8265.pem",
|
|
}
|
|
|
|
|
|
def test_create_node_returns_dict():
|
|
mock_node_config = {"machineType": "n2-standard-8"}
|
|
mock_results = [({"dict": 1}, "instance_id1"), ({"dict": 2}, "instance_id2")]
|
|
mock_resource = MagicMock()
|
|
mock_resource.create_instances.return_value = mock_results
|
|
expected_return_value = {"instance_id1": {"dict": 1}, "instance_id2": {"dict": 2}}
|
|
|
|
def __init__(self, provider_config: dict, cluster_name: str):
|
|
self.lock = RLock()
|
|
self.cached_nodes: Dict[str, GCPNode] = {}
|
|
self.resources: Dict[GCPNodeType, GCPResource] = {}
|
|
self.cache_stopped_nodes = False
|
|
self.resources[GCPNodeType.COMPUTE] = mock_resource
|
|
|
|
with patch.object(GCPNodeProvider, "__init__", __init__):
|
|
node_provider = GCPNodeProvider({}, "")
|
|
create_node_return_value = node_provider.create_node(mock_node_config, {}, 1)
|
|
assert create_node_return_value == expected_return_value
|
|
|
|
|
|
def test_terminate_nodes():
|
|
mock_node_config = {"machineType": "n2-standard-8"}
|
|
node_type = GCPNodeType.COMPUTE.value
|
|
id1, id2 = f"instance-id1-{node_type}", f"instance-id2-{node_type}"
|
|
terminate_node_ids = [id1, id2]
|
|
mock_resource = MagicMock()
|
|
mock_resource.create_instances.return_value = [
|
|
({"dict": 1}, id1),
|
|
({"dict": 2}, id2),
|
|
]
|
|
mock_resource.delete_instance.return_value = "test"
|
|
|
|
def __init__(self, provider_config: dict, cluster_name: str):
|
|
self.lock = RLock()
|
|
self.cached_nodes: Dict[str, GCPNode] = {}
|
|
self.resources: Dict[GCPNodeType, GCPResource] = {}
|
|
self.cache_stopped_nodes = False
|
|
self.resources[GCPNodeType.COMPUTE] = mock_resource
|
|
|
|
with patch.object(GCPNodeProvider, "__init__", __init__):
|
|
node_provider = GCPNodeProvider({}, "")
|
|
node_provider.create_node(mock_node_config, {}, 1)
|
|
node_provider.terminate_nodes(terminate_node_ids)
|
|
|
|
mock_resource.delete_instance.assert_has_calls(
|
|
[call(node_id=id1), call(node_id=id2)], any_order=True
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
("n1-standard-4", f"zones/{_AZ}/machineTypes/n1-standard-4"),
|
|
(
|
|
f"zones/{_AZ}/machineTypes/n1-standard-4",
|
|
f"zones/{_AZ}/machineTypes/n1-standard-4",
|
|
),
|
|
],
|
|
)
|
|
def test_convert_resources_to_urls_machine(test_case):
|
|
gcp_compute = GCPCompute(None, _PROJECT_NAME, _AZ, "cluster_name")
|
|
base_machine, result_machine = test_case
|
|
modified_config = gcp_compute._convert_resources_to_urls(
|
|
{"machineType": base_machine}
|
|
)
|
|
assert modified_config["machineType"] == result_machine
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
(
|
|
"nvidia-tesla-k80",
|
|
f"projects/{_PROJECT_NAME}/zones/{_AZ}/acceleratorTypes/nvidia-tesla-k80",
|
|
),
|
|
(
|
|
f"projects/{_PROJECT_NAME}/zones/{_AZ}/acceleratorTypes/nvidia-tesla-k80",
|
|
f"projects/{_PROJECT_NAME}/zones/{_AZ}/acceleratorTypes/nvidia-tesla-k80",
|
|
),
|
|
],
|
|
)
|
|
def test_convert_resources_to_urls_accelerators(test_case):
|
|
gcp_compute = GCPCompute(None, _PROJECT_NAME, _AZ, "cluster_name")
|
|
base_accel, result_accel = test_case
|
|
|
|
base_config = {
|
|
"machineType": "n1-standard-4",
|
|
"guestAccelerators": [{"acceleratorCount": 1, "acceleratorType": base_accel}],
|
|
}
|
|
modified_config = gcp_compute._convert_resources_to_urls(base_config)
|
|
|
|
assert modified_config["guestAccelerators"][0]["acceleratorType"] == result_accel
|
|
|
|
|
|
def test_compute_node_list_instances_excludes_tpu():
|
|
mock_execute = MagicMock(return_value={"test": "abc"})
|
|
mock_list = MagicMock(return_value=MagicMock(execute=mock_execute))
|
|
mock_instances = MagicMock(return_value=MagicMock(list=mock_list))
|
|
mock_resource = MagicMock(instances=mock_instances)
|
|
|
|
GCPCompute(mock_resource, _PROJECT_NAME, _AZ, "cluster_name").list_instances()
|
|
filter_arg = mock_list.call_args.kwargs["filter"]
|
|
|
|
# Checks that the tpu negation filter is included.
|
|
assert "tpu_cores" in filter_arg
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
(
|
|
{},
|
|
SSHCommandRunner,
|
|
),
|
|
(
|
|
{"docker_config": {"container_name": "container"}},
|
|
DockerCommandRunner,
|
|
),
|
|
],
|
|
)
|
|
def test_cpu_resource_returns_standard_command_runner(test_case):
|
|
mock_resource = MagicMock()
|
|
|
|
def __init__(self, provider_config: dict, cluster_name: str):
|
|
self.lock = RLock()
|
|
self.cached_nodes: Dict[str, GCPNode] = {}
|
|
self.resources: Dict[GCPNodeType, GCPResource] = {}
|
|
self.resources[GCPNodeType.COMPUTE] = mock_resource
|
|
|
|
with patch.object(GCPNodeProvider, "__init__", __init__):
|
|
node_provider = GCPNodeProvider({}, "")
|
|
|
|
optional_docker_config, expected_runner = test_case
|
|
|
|
args = {
|
|
"log_prefix": "test",
|
|
"node_id": "test-instance-compute",
|
|
"auth_config": auth_config,
|
|
"cluster_name": "test",
|
|
"process_runner": MockProcessRunner(),
|
|
"use_internal_ip": True,
|
|
}
|
|
args.update(optional_docker_config)
|
|
command_runner = node_provider.get_command_runner(**args)
|
|
assert isinstance(command_runner, expected_runner)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
(
|
|
{},
|
|
TPUVMSSHCommandRunner,
|
|
),
|
|
(
|
|
{"docker_config": {"container_name": "container"}},
|
|
TPUVMDockerCommandRunner,
|
|
),
|
|
],
|
|
)
|
|
def test_tpu_resource_returns_tpu_command_runner(test_case):
|
|
mock_resource = MagicMock()
|
|
|
|
def __init__(self, provider_config: dict, cluster_name: str):
|
|
self.lock = RLock()
|
|
self.cached_nodes: Dict[str, GCPNode] = {}
|
|
self.resources: Dict[GCPNodeType, GCPResource] = {}
|
|
self.resources[GCPNodeType.COMPUTE] = mock_resource
|
|
self.resources[GCPNodeType.TPU] = mock_resource
|
|
|
|
with patch.object(GCPNodeProvider, "__init__", __init__):
|
|
node_provider = GCPNodeProvider({}, "")
|
|
|
|
optional_docker_config, expected_runner = test_case
|
|
|
|
args = {
|
|
"log_prefix": "test",
|
|
"node_id": "test-instance-tpu",
|
|
"auth_config": auth_config,
|
|
"cluster_name": "test",
|
|
"process_runner": MockProcessRunner(),
|
|
"use_internal_ip": True,
|
|
}
|
|
args.update(optional_docker_config)
|
|
command_runner = node_provider.get_command_runner(**args)
|
|
assert isinstance(command_runner, TPUCommandRunner)
|
|
assert isinstance(command_runner._command_runners[0], expected_runner)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
({"acceleratorType": "v4-16"}, "TPU-v4-16-head"),
|
|
({"acceleratorType": "v4-32"}, "TPU-v4-32-head"),
|
|
({"acceleratorType": "v3-8"}, "TPU-v3-8-head"),
|
|
({"acceleratorConfig": {"type": "V4", "topology": "2x2x2"}}, "TPU-v4-16-head"),
|
|
({"acceleratorConfig": {"type": "V4", "topology": "4x4x4"}}, "TPU-v4-128-head"),
|
|
(
|
|
{"acceleratorConfig": {"type": "V5LITE_POD", "topology": "2x4"}},
|
|
"TPU-v5litepod-8-head",
|
|
),
|
|
(
|
|
{"acceleratorConfig": {"type": "V6E", "topology": "2x4"}},
|
|
"TPU-v6e-8-head",
|
|
),
|
|
],
|
|
)
|
|
def test_tpu_node_fillout(test_case):
|
|
accelerator_config, expected_resource_str = test_case
|
|
|
|
cluster_config = {
|
|
"available_node_types": {
|
|
"ray_tpu": {
|
|
"resources": {"TPU": 4},
|
|
"node_config": {
|
|
"runtimeVersion": "tpu-vm-v4-base",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
cluster_config["available_node_types"]["ray_tpu"]["node_config"].update(
|
|
accelerator_config
|
|
)
|
|
|
|
new_config = GCPNodeProvider.fillout_available_node_types_resources(
|
|
cluster_config=cluster_config
|
|
)
|
|
resource_config = new_config["available_node_types"]["ray_tpu"]["resources"]
|
|
assert expected_resource_str in resource_config
|
|
assert resource_config[expected_resource_str] == 1
|
|
|
|
|
|
def test_tpu_config_cannot_have_accelerator_type_and_config():
|
|
node = {
|
|
"acceleratorType": "abc",
|
|
"acceleratorConfig": {"abc": "def"},
|
|
}
|
|
with pytest.raises(ValueError):
|
|
get_node_type(node)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"node",
|
|
[
|
|
{"acceleratorConfig": {"type": "V3", "topology": "2x2"}},
|
|
{"acceleratorConfig": {"type": "V2", "topology": "2x2"}},
|
|
],
|
|
)
|
|
def test_get_node_rejects_v2_v3_accelerator_config(node):
|
|
with pytest.raises(ValueError):
|
|
get_node_type(node)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"node_config",
|
|
[
|
|
{"acceleratorType": "vabc-12345"},
|
|
{"acceleratorType": "v3-abc"},
|
|
{"acceleratorType": "v3-8a"},
|
|
{"acceleratorType": "this should fail"},
|
|
{"acceleratorConfig": {"type": "asdf", "topology": "2x2x1"}},
|
|
{"acceleratorConfig": {"type": "V4", "topology": "asdf"}},
|
|
],
|
|
)
|
|
def test_invalid_accelerator_configs(node_config):
|
|
with pytest.raises(ValueError):
|
|
get_node_type(node_config)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
({"acceleratorType": "v2-8"}, 4, True),
|
|
({"acceleratorType": "v3-8"}, 4, True),
|
|
({"acceleratorType": "v4-8"}, 4, True),
|
|
# Note: Topology only supported in v4
|
|
({"acceleratorConfig": {"type": "V4", "topology": "2x2x1"}}, 4, True),
|
|
({"acceleratorType": "v2-32"}, 16, False),
|
|
({"acceleratorType": "v3-128"}, 64, False),
|
|
({"acceleratorType": "v4-4096"}, 2048, False),
|
|
({"acceleratorConfig": {"type": "V4", "topology": "2x2x8"}}, 32, False),
|
|
({"acceleratorConfig": {"type": "V4", "topology": "4x4x4"}}, 64, False),
|
|
({"acceleratorConfig": {"type": "V5LITE_POD", "topology": "2x4"}}, 8, True),
|
|
({"acceleratorConfig": {"type": "V6E", "topology": "2x4"}}, 8, True),
|
|
],
|
|
)
|
|
def test_tpu_chip_calculation_single_host_logic(test_case):
|
|
node, expected_chips, expected_singlehost = test_case
|
|
assert _get_num_tpu_chips(node) == expected_chips
|
|
assert _is_single_host_tpu(node) == expected_singlehost
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
({"machineType": "n2-standard-4"}, GCPNodeType.COMPUTE, False),
|
|
(
|
|
{
|
|
"machineType": "n2-standard-4",
|
|
"acceleratorType": {
|
|
"guestAccelerators": {
|
|
"acceleratorType": "V100",
|
|
"acceleratorCount": 1,
|
|
}
|
|
},
|
|
},
|
|
GCPNodeType.COMPUTE,
|
|
False,
|
|
),
|
|
({"acceleratorType": "v2-8"}, GCPNodeType.TPU, True),
|
|
({"acceleratorType": "v3-8"}, GCPNodeType.TPU, True),
|
|
({"acceleratorType": "v4-8"}, GCPNodeType.TPU, True),
|
|
(
|
|
{"acceleratorConfig": {"type": "V4", "topology": "2x2x1"}},
|
|
GCPNodeType.TPU,
|
|
True,
|
|
),
|
|
({"acceleratorType": "v2-32"}, GCPNodeType.TPU, True),
|
|
({"acceleratorType": "v3-128"}, GCPNodeType.TPU, True),
|
|
({"acceleratorType": "v4-4096"}, GCPNodeType.TPU, True),
|
|
(
|
|
{"acceleratorConfig": {"type": "V4", "topology": "2x2x8"}},
|
|
GCPNodeType.TPU,
|
|
True,
|
|
),
|
|
(
|
|
{"acceleratorConfig": {"type": "V4", "topology": "4x4x4"}},
|
|
GCPNodeType.TPU,
|
|
True,
|
|
),
|
|
(
|
|
{"acceleratorConfig": {"type": "V5LITE_POD", "topology": "2x4"}},
|
|
GCPNodeType.TPU,
|
|
True,
|
|
),
|
|
(
|
|
{"acceleratorConfig": {"type": "V6E", "topology": "2x4"}},
|
|
GCPNodeType.TPU,
|
|
True,
|
|
),
|
|
],
|
|
)
|
|
def test_get_node_type_and_has_tpu(test_case):
|
|
node, expected_compute_type, expected_is_tpu = test_case
|
|
assert get_node_type(node) == expected_compute_type
|
|
config = {
|
|
"available_node_types": {
|
|
"node_type_1": {"node_config": node},
|
|
},
|
|
}
|
|
assert _has_tpus_in_node_configs(config) == expected_is_tpu
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"accelerator_pod_tuple",
|
|
[
|
|
({"acceleratorType": "v2-32"}, True),
|
|
({"acceleratorType": "v3-32"}, True),
|
|
({"acceleratorType": "v4-32"}, True),
|
|
({"acceleratorConfig": {"type": "V4", "topology": "2x2x2"}}, True),
|
|
({"acceleratorType": "v2-8"}, False),
|
|
({"acceleratorType": "v3-8"}, False),
|
|
({"acceleratorType": "v4-8"}, False),
|
|
({"acceleratorConfig": {"type": "V4", "topology": "2x2x1"}}, False),
|
|
],
|
|
)
|
|
def test_tpu_pod_emits_warning(propagate_logs, caplog, accelerator_pod_tuple):
|
|
accelerator, should_emit = accelerator_pod_tuple
|
|
|
|
with caplog.at_level(
|
|
logging.WARNING, logger="ray.autoscaler._private.gcp.config.get_node_type"
|
|
):
|
|
get_node_type(accelerator)
|
|
if should_emit:
|
|
assert "TPU pod detected" in caplog.text
|
|
else:
|
|
assert "TPU pod detected" not in caplog.text
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
("v4-8", "V4", "2x2x1"),
|
|
("v4-16", "V4", "2x2x2"),
|
|
("v4-128", "V4", "4x4x4"),
|
|
("v4-256", "V4", "4x4x8"),
|
|
("v5litepod-8", "V5LITE_POD", "2x4"),
|
|
("v6e-8", "V6E", "2x4"),
|
|
],
|
|
)
|
|
def test_tpu_accelerator_config_to_type(test_case):
|
|
expected, accel_type, topology = test_case
|
|
accelerator_config = {
|
|
"type": accel_type,
|
|
"topology": topology,
|
|
}
|
|
accelerator_type = tpu_accelerator_config_to_type(
|
|
accelerator_config=accelerator_config
|
|
)
|
|
assert accelerator_type == expected
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|