wehub-resource-sync
234 lines
7.5 KiB
Python
234 lines
7.5 KiB
Python
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray._private.accelerators.tpu import TPUAcceleratorManager
|
|
from ray.cluster_utils import AutoscalingCluster
|
|
|
|
|
|
def check_cmd_stderr(cmd):
|
|
return subprocess.run(cmd, stderr=subprocess.PIPE).stderr.decode("utf-8")
|
|
|
|
|
|
def add_default_labels_for_test(node_info, labels):
|
|
labels["ray.io/node-id"] = node_info["NodeID"]
|
|
return labels
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"call_ray_start",
|
|
['ray start --head --labels={"gpu_type":"A100","region":"us"}'],
|
|
indirect=True,
|
|
)
|
|
def test_ray_start_set_node_labels_from_json(call_ray_start):
|
|
ray.init(address=call_ray_start)
|
|
node_info = ray.nodes()[0]
|
|
assert node_info["Labels"] == add_default_labels_for_test(
|
|
node_info, {"gpu_type": "A100", "region": "us"}
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"call_ray_start",
|
|
['ray start --head --labels "gpu_type=A100,region=us"'],
|
|
indirect=True,
|
|
)
|
|
def test_ray_start_set_node_labels_from_string(call_ray_start):
|
|
ray.init(address=call_ray_start)
|
|
node_info = ray.nodes()[0]
|
|
assert node_info["Labels"] == add_default_labels_for_test(
|
|
node_info, {"gpu_type": "A100", "region": "us"}
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"call_ray_start",
|
|
[
|
|
"ray start --head --labels={}",
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_ray_start_set_empty_node_labels(call_ray_start):
|
|
ray.init(address=call_ray_start)
|
|
node_info = ray.nodes()[0]
|
|
assert node_info["Labels"] == add_default_labels_for_test(node_info, {})
|
|
|
|
|
|
def test_ray_init_set_node_labels(shutdown_only):
|
|
labels = {"gpu_type": "A100", "region": "us"}
|
|
ray.init(labels=labels)
|
|
node_info = ray.nodes()[0]
|
|
assert node_info["Labels"] == add_default_labels_for_test(node_info, labels)
|
|
ray.shutdown()
|
|
ray.init(labels={})
|
|
node_info = ray.nodes()[0]
|
|
assert node_info["Labels"] == add_default_labels_for_test(node_info, {})
|
|
|
|
|
|
def test_node_labels_connect_only_false(shutdown_only):
|
|
labels = {"key": "value"}
|
|
ray.init(labels=labels)
|
|
node_labels = ray.get_runtime_context().get_node_labels()
|
|
assert node_labels["key"] == "value"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"call_ray_start",
|
|
['ray start --head --labels={"key":"value"}'],
|
|
indirect=True,
|
|
)
|
|
def test_node_labels_connect_only_true(call_ray_start):
|
|
ray.init(address=call_ray_start)
|
|
node_labels = ray.get_runtime_context().get_node_labels()
|
|
assert node_labels["key"] == "value"
|
|
|
|
|
|
def test_ray_init_set_node_labels_value_error(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
|
|
cluster.add_node(num_cpus=1)
|
|
with pytest.raises(ValueError, match="labels must not be provided"):
|
|
ray.init(address=cluster.address, labels={"gpu_type": "A100"})
|
|
|
|
with pytest.raises(ValueError, match="labels must not be provided"):
|
|
ray.init(labels={"gpu_type": "A100"})
|
|
|
|
|
|
def test_ray_start_set_node_labels_value_error():
|
|
out = check_cmd_stderr(["ray", "start", "--head", "--labels=xxx"])
|
|
assert "Label string is not a key-value pair." in out
|
|
|
|
out = check_cmd_stderr(["ray", "start", "--head", '--labels={"gpu_type":1}'])
|
|
assert "Label string is not a key-value pair." in out
|
|
|
|
out = check_cmd_stderr(
|
|
["ray", "start", "--head", '--labels={"ray.io/node-id":"111"}']
|
|
)
|
|
assert "Label string is not a key-value pair" in out
|
|
|
|
out = check_cmd_stderr(
|
|
["ray", "start", "--head", '--labels={"ray.io/other_key":"111"}']
|
|
)
|
|
assert "Label string is not a key-value pair" in out
|
|
|
|
|
|
def test_cluster_add_node_with_labels(ray_start_cluster):
|
|
labels = {"gpu_type": "A100", "region": "us"}
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=1, labels=labels)
|
|
cluster.wait_for_nodes()
|
|
ray.init(address=cluster.address)
|
|
node_info = ray.nodes()[0]
|
|
assert node_info["Labels"] == add_default_labels_for_test(node_info, labels)
|
|
head_node_id = ray.nodes()[0]["NodeID"]
|
|
|
|
cluster.add_node(num_cpus=1, labels={})
|
|
cluster.wait_for_nodes()
|
|
for node in ray.nodes():
|
|
if node["NodeID"] != head_node_id:
|
|
assert node["Labels"] == add_default_labels_for_test(node, {})
|
|
|
|
|
|
@pytest.mark.parametrize("autoscaler_v2", [False, True], ids=["v1", "v2"])
|
|
def test_autoscaler_set_node_labels(autoscaler_v2, shutdown_only):
|
|
cluster = AutoscalingCluster(
|
|
head_resources={"CPU": 0},
|
|
worker_node_types={
|
|
"worker_1": {
|
|
"resources": {"CPU": 1},
|
|
"labels": {"region": "us"},
|
|
"node_config": {},
|
|
"min_workers": 1,
|
|
"max_workers": 1,
|
|
}
|
|
},
|
|
autoscaler_v2=autoscaler_v2,
|
|
)
|
|
|
|
try:
|
|
cluster.start()
|
|
ray.init()
|
|
wait_for_condition(lambda: len(ray.nodes()) == 2, timeout=20)
|
|
|
|
for node in ray.nodes():
|
|
if node["Resources"].get("CPU", 0) == 1:
|
|
assert node["Labels"] == add_default_labels_for_test(
|
|
node, {"region": "us"}
|
|
)
|
|
finally:
|
|
cluster.shutdown()
|
|
|
|
|
|
def test_ray_start_set_node_labels_from_file(shutdown_only):
|
|
with tempfile.NamedTemporaryFile(mode="w+", delete=False) as test_file:
|
|
test_file.write('"gpu_type": "A100"\n"region": "us"\n"market-type": "spot"')
|
|
test_file_path = test_file.name
|
|
|
|
try:
|
|
cmd = ["ray", "start", "--head", "--labels-file", test_file_path]
|
|
subprocess.check_call(cmd)
|
|
ray.init(address="auto")
|
|
node_info = ray.nodes()[0]
|
|
assert node_info["Labels"] == add_default_labels_for_test(
|
|
node_info, {"gpu_type": "A100", "region": "us", "market-type": "spot"}
|
|
)
|
|
finally:
|
|
subprocess.check_call(["ray", "stop", "--force"])
|
|
os.remove(test_file_path)
|
|
|
|
|
|
def test_get_default_ray_node_labels(shutdown_only, monkeypatch):
|
|
# Set env vars for this test
|
|
monkeypatch.setenv("RAY_NODE_MARKET_TYPE", "spot")
|
|
monkeypatch.setenv("RAY_NODE_TYPE_NAME", "worker-group-1")
|
|
monkeypatch.setenv("RAY_NODE_REGION", "us-central2")
|
|
monkeypatch.setenv("RAY_NODE_ZONE", "us-central2-b")
|
|
monkeypatch.setenv("TPU_ACCELERATOR_TYPE", "v4-16")
|
|
|
|
ray.init(resources={"TPU": 4})
|
|
node_info = ray.nodes()[0]
|
|
labels = node_info["Labels"]
|
|
|
|
assert labels.get("ray.io/market-type") == "spot"
|
|
assert labels.get("ray.io/node-group") == "worker-group-1"
|
|
assert labels.get("ray.io/availability-region") == "us-central2"
|
|
assert labels.get("ray.io/availability-zone") == "us-central2-b"
|
|
assert labels.get("ray.io/accelerator-type") == "TPU-V4"
|
|
|
|
|
|
def test_get_default_tpu_labels(shutdown_only, monkeypatch):
|
|
# Set env vars for this test
|
|
monkeypatch.setenv("TPU_NAME", "slice-0")
|
|
monkeypatch.setenv("TPU_WORKER_ID", "0")
|
|
monkeypatch.setenv("TPU_ACCELERATOR_TYPE", "v6e-32")
|
|
monkeypatch.setenv("TPU_TOPOLOGY", "4x8")
|
|
|
|
with patch(
|
|
"ray._private.accelerators.get_all_accelerator_resource_names",
|
|
return_value=["TPU"],
|
|
), patch(
|
|
"ray._private.accelerators.get_accelerator_manager_for_resource",
|
|
return_value=TPUAcceleratorManager(),
|
|
):
|
|
ray.init(resources={"TPU": 4})
|
|
node_info = ray.nodes()[0]
|
|
labels = node_info["Labels"]
|
|
|
|
assert labels.get("ray.io/accelerator-type") == "TPU-V6E"
|
|
|
|
# TPU specific labels for SPMD
|
|
assert labels.get("ray.io/tpu-slice-name") == "slice-0"
|
|
assert labels.get("ray.io/tpu-worker-id") == "0"
|
|
assert labels.get("ray.io/tpu-topology") == "4x8"
|
|
assert labels.get("ray.io/tpu-pod-type") == "v6e-32"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|