chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,444 @@
|
||||
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__]))
|
||||
@@ -0,0 +1,290 @@
|
||||
import hashlib
|
||||
import os
|
||||
import sys
|
||||
from getpass import getuser
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from ray._private import ray_constants
|
||||
from ray.autoscaler._private.command_runner import SSHCommandRunner
|
||||
from ray.autoscaler._private.gcp.tpu_command_runner import TPUCommandRunner
|
||||
from ray.tests.test_autoscaler import MockProcessRunner, MockProvider
|
||||
|
||||
_MOCK_TPU_NAME = "my-tpu"
|
||||
_MOCK_ACCELERATOR_TYPE = "v4-16"
|
||||
|
||||
auth_config = {
|
||||
"ssh_user": "ray",
|
||||
"ssh_private_key": "8265.pem",
|
||||
}
|
||||
|
||||
|
||||
class MockTpuInstance:
|
||||
def __init__(self, num_workers: int = 1):
|
||||
self.num_workers = num_workers
|
||||
|
||||
def get_internal_ip(self, worker_index: int) -> str:
|
||||
return "0.0.0.0"
|
||||
|
||||
def get_external_ip(self, worker_index: int) -> str:
|
||||
return "1.2.3.4"
|
||||
|
||||
def get(self, key) -> str:
|
||||
if key == "name":
|
||||
return _MOCK_TPU_NAME
|
||||
elif key == "acceleratorType":
|
||||
return _MOCK_ACCELERATOR_TYPE
|
||||
return ""
|
||||
|
||||
|
||||
def test_tpu_ssh_command_runner():
|
||||
num_workers = 2
|
||||
process_runner = MockProcessRunner()
|
||||
provider = MockProvider()
|
||||
instance = MockTpuInstance(num_workers=num_workers)
|
||||
provider.create_node({}, {}, 1)
|
||||
cluster_name = "cluster"
|
||||
ssh_control_hash = hashlib.sha256(cluster_name.encode()).hexdigest()
|
||||
ssh_user_hash = hashlib.sha256(getuser().encode()).hexdigest()
|
||||
ssh_control_path = "/tmp/ray_ssh_{}/{}".format(
|
||||
ssh_user_hash[:10], ssh_control_hash[:10]
|
||||
)
|
||||
args = {
|
||||
"instance": instance,
|
||||
"log_prefix": "prefix",
|
||||
"node_id": "abc",
|
||||
"provider": provider,
|
||||
"auth_config": auth_config,
|
||||
"cluster_name": cluster_name,
|
||||
"process_runner": process_runner,
|
||||
"use_internal_ip": False,
|
||||
}
|
||||
env_vars = {"var1": 'quote between this " and this', "var2": "123"}
|
||||
cmd_runner = TPUCommandRunner(**args)
|
||||
cmd_runner.run(
|
||||
"echo helloo", port_forward=[(8265, 8265)], environment_variables=env_vars
|
||||
)
|
||||
|
||||
expected = [
|
||||
"ssh",
|
||||
"-tt",
|
||||
"-L",
|
||||
"8265:localhost:8265",
|
||||
"-i",
|
||||
"8265.pem",
|
||||
"-o",
|
||||
"StrictHostKeyChecking=no",
|
||||
"-o",
|
||||
"UserKnownHostsFile=/dev/null",
|
||||
"-o",
|
||||
"IdentitiesOnly=yes",
|
||||
"-o",
|
||||
"ExitOnForwardFailure=yes",
|
||||
"-o",
|
||||
"ServerAliveInterval=5",
|
||||
"-o",
|
||||
"ServerAliveCountMax=3",
|
||||
"-o",
|
||||
"ControlMaster=auto",
|
||||
"-o",
|
||||
"ControlPath={}/%C".format(ssh_control_path),
|
||||
"-o",
|
||||
"ControlPersist=10s",
|
||||
"-o",
|
||||
"ConnectTimeout=120s",
|
||||
"ray@1.2.3.4",
|
||||
"bash",
|
||||
"--login",
|
||||
"-c",
|
||||
"-i",
|
||||
"""'source ~/.bashrc; export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && (export var1='"'"'"quote between this \\" and this"'"'"';export var2='"'"'"123"'"'"';echo helloo)'""", # noqa: E501
|
||||
]
|
||||
|
||||
calls = process_runner.calls
|
||||
|
||||
# Asserts that we do make the call once per worker in the TPU pod.
|
||||
assert len(process_runner.calls) == num_workers
|
||||
|
||||
# Much easier to debug this loop than the function call.
|
||||
for i in range(num_workers):
|
||||
for x, y in zip(calls[i], expected):
|
||||
assert x == y
|
||||
|
||||
|
||||
def test_tpu_docker_command_runner():
|
||||
num_workers = 4
|
||||
process_runner = MockProcessRunner()
|
||||
provider = MockProvider()
|
||||
instance = MockTpuInstance(num_workers=num_workers)
|
||||
provider.create_node({}, {}, 1)
|
||||
cluster_name = "cluster"
|
||||
ssh_control_hash = hashlib.sha256(cluster_name.encode()).hexdigest()
|
||||
ssh_user_hash = hashlib.sha256(getuser().encode()).hexdigest()
|
||||
ssh_control_path = "/tmp/ray_ssh_{}/{}".format(
|
||||
ssh_user_hash[:10], ssh_control_hash[:10]
|
||||
)
|
||||
docker_config = {"container_name": "container"}
|
||||
args = {
|
||||
"instance": instance,
|
||||
"log_prefix": "prefix",
|
||||
"node_id": "0",
|
||||
"provider": provider,
|
||||
"auth_config": auth_config,
|
||||
"cluster_name": cluster_name,
|
||||
"process_runner": process_runner,
|
||||
"use_internal_ip": False,
|
||||
"docker_config": docker_config,
|
||||
}
|
||||
cmd_runner = TPUCommandRunner(**args)
|
||||
env_vars = {"var1": 'quote between this " and this', "var2": "123"}
|
||||
cmd_runner.run("echo hello", environment_variables=env_vars)
|
||||
|
||||
# This string is insane because there are an absurd number of embedded
|
||||
# quotes. While this is a ridiculous string, the escape behavior is
|
||||
# important and somewhat difficult to get right for environment variables.
|
||||
cmd = """'source ~/.bashrc; export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && (docker exec -it container /bin/bash -c '"'"'bash --login -c -i '"'"'"'"'"'"'"'"'source ~/.bashrc; export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && (export var1='"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"quote between this \\" and this"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"';export var2='"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"123"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"';echo hello)'"'"'"'"'"'"'"'"''"'"' )'""" # noqa: E501
|
||||
|
||||
expected = [
|
||||
"ssh",
|
||||
"-tt",
|
||||
"-i",
|
||||
"8265.pem",
|
||||
"-o",
|
||||
"StrictHostKeyChecking=no",
|
||||
"-o",
|
||||
"UserKnownHostsFile=/dev/null",
|
||||
"-o",
|
||||
"IdentitiesOnly=yes",
|
||||
"-o",
|
||||
"ExitOnForwardFailure=yes",
|
||||
"-o",
|
||||
"ServerAliveInterval=5",
|
||||
"-o",
|
||||
"ServerAliveCountMax=3",
|
||||
"-o",
|
||||
"ControlMaster=auto",
|
||||
"-o",
|
||||
"ControlPath={}/%C".format(ssh_control_path),
|
||||
"-o",
|
||||
"ControlPersist=10s",
|
||||
"-o",
|
||||
"ConnectTimeout=120s",
|
||||
"ray@1.2.3.4",
|
||||
"bash",
|
||||
"--login",
|
||||
"-c",
|
||||
"-i",
|
||||
cmd,
|
||||
]
|
||||
|
||||
calls = process_runner.calls
|
||||
|
||||
# Asserts that we do make the call once per worker in the TPU pod.
|
||||
assert len(process_runner.calls) == num_workers
|
||||
|
||||
# Much easier to debug this loop than the function call.
|
||||
for i in range(num_workers):
|
||||
for x, y in zip(calls[i], expected):
|
||||
assert x == y
|
||||
|
||||
|
||||
def test_tpu_docker_run_init():
|
||||
num_workers = 1
|
||||
process_runner = MockProcessRunner()
|
||||
provider = MockProvider()
|
||||
instance = MockTpuInstance(num_workers=num_workers)
|
||||
provider.create_node({}, {}, 1)
|
||||
cluster_name = "cluster"
|
||||
docker_config = {
|
||||
"container_name": "container",
|
||||
"image": "rayproject/ray:latest",
|
||||
}
|
||||
args = {
|
||||
"instance": instance,
|
||||
"log_prefix": "prefix",
|
||||
"node_id": "0",
|
||||
"provider": provider,
|
||||
"auth_config": auth_config,
|
||||
"cluster_name": cluster_name,
|
||||
"process_runner": process_runner,
|
||||
"use_internal_ip": False,
|
||||
"docker_config": docker_config,
|
||||
}
|
||||
cmd_runner = TPUCommandRunner(**args)
|
||||
|
||||
# Taken from tests/test_command_runner.py
|
||||
# This mocks the response of 'docker inspect' command to return an empty JSON array.
|
||||
# This simulates the scenario where the Docker image has no set environment
|
||||
# variables, allowing us to test the subsequent code for handling this case.
|
||||
process_runner.respond_to_call("json .Config.Env", 2 * ["[]"])
|
||||
cmd_runner.run_init(as_head=True, file_mounts={}, sync_run_yet=True)
|
||||
process_runner.assert_has_call("1.2.3.4", pattern="docker")
|
||||
|
||||
|
||||
def test_max_active_connections_env_var():
|
||||
num_workers = 2
|
||||
process_runner = MockProcessRunner()
|
||||
provider = MockProvider()
|
||||
instance = MockTpuInstance(num_workers=num_workers)
|
||||
provider.create_node({}, {}, 1)
|
||||
cluster_name = "cluster"
|
||||
docker_config = {"container_name": "container"}
|
||||
args = {
|
||||
"instance": instance,
|
||||
"log_prefix": "prefix",
|
||||
"node_id": "0",
|
||||
"provider": provider,
|
||||
"auth_config": auth_config,
|
||||
"cluster_name": cluster_name,
|
||||
"process_runner": process_runner,
|
||||
"use_internal_ip": False,
|
||||
"docker_config": docker_config,
|
||||
}
|
||||
cmd_runner = TPUCommandRunner(**args)
|
||||
os.environ[ray_constants.RAY_TPU_MAX_CONCURRENT_CONNECTIONS_ENV_VAR] = "1"
|
||||
num_connections = cmd_runner.num_connections
|
||||
assert type(num_connections) is int
|
||||
assert num_connections == 1
|
||||
|
||||
|
||||
def test_tpu_pod_resources():
|
||||
num_workers = 2
|
||||
process_runner = MockProcessRunner()
|
||||
provider = MockProvider()
|
||||
instance = MockTpuInstance(num_workers=num_workers)
|
||||
provider.create_node({}, {}, 1)
|
||||
cluster_name = "cluster"
|
||||
args = {
|
||||
"instance": instance,
|
||||
"log_prefix": "prefix",
|
||||
"node_id": "abc",
|
||||
"provider": provider,
|
||||
"auth_config": auth_config,
|
||||
"cluster_name": cluster_name,
|
||||
"process_runner": process_runner,
|
||||
"use_internal_ip": False,
|
||||
}
|
||||
env_vars = {
|
||||
ray_constants.RESOURCES_ENVIRONMENT_VARIABLE: {
|
||||
"TPU": 4,
|
||||
f"TPU-{_MOCK_ACCELERATOR_TYPE}-head": 1,
|
||||
},
|
||||
}
|
||||
|
||||
def test_command_run(self, environment_variables, **kwargs):
|
||||
resources = environment_variables[ray_constants.RESOURCES_ENVIRONMENT_VARIABLE]
|
||||
if self._worker_id == 0:
|
||||
assert f"TPU-{_MOCK_ACCELERATOR_TYPE}-head" in resources
|
||||
else:
|
||||
assert f"TPU-{_MOCK_ACCELERATOR_TYPE}-head" not in resources
|
||||
|
||||
with patch.object(SSHCommandRunner, "run", new=test_command_run):
|
||||
cmd_runner = TPUCommandRunner(**args)
|
||||
cmd_runner.run(
|
||||
"echo helloo", port_forward=[(8265, 8265)], environment_variables=env_vars
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
Reference in New Issue
Block a user