wehub-resource-sync
1025 lines
32 KiB
Python
1025 lines
32 KiB
Python
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from typing import Dict
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.constants import HEAD_NODE_RESOURCE_NAME
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray._common.usage.usage_lib import get_extra_usage_tags_to_report
|
|
from ray._private.test_utils import run_string_as_driver_nonblocking
|
|
from ray._raylet import GcsClient
|
|
from ray.autoscaler.v2.sdk import get_cluster_status
|
|
from ray.cluster_utils import AutoscalingCluster
|
|
from ray.core.generated.usage_pb2 import TagKey
|
|
from ray.util.placement_group import (
|
|
placement_group,
|
|
remove_placement_group,
|
|
)
|
|
from ray.util.state.api import (
|
|
list_actors,
|
|
list_placement_groups,
|
|
list_tasks,
|
|
)
|
|
|
|
|
|
def is_head_node_from_resource_usage(usage: Dict[str, float]) -> bool:
|
|
if HEAD_NODE_RESOURCE_NAME in usage:
|
|
return True
|
|
return False
|
|
|
|
|
|
@pytest.mark.parametrize("autoscaler_v2", [False, True], ids=["v1", "v2"])
|
|
def test_autoscaler_no_churn(autoscaler_v2):
|
|
num_cpus_per_node = 4
|
|
expected_nodes = 6
|
|
cluster = AutoscalingCluster(
|
|
head_resources={"CPU": num_cpus_per_node},
|
|
worker_node_types={
|
|
"type-1": {
|
|
"resources": {"CPU": num_cpus_per_node},
|
|
"node_config": {},
|
|
"min_workers": 0,
|
|
"max_workers": 2 * expected_nodes,
|
|
},
|
|
},
|
|
autoscaler_v2=autoscaler_v2,
|
|
)
|
|
|
|
driver_script = f"""
|
|
import time
|
|
import ray
|
|
@ray.remote(num_cpus=1)
|
|
def foo():
|
|
time.sleep(60)
|
|
return True
|
|
|
|
ray.init("auto")
|
|
|
|
print("start")
|
|
assert(ray.get([foo.remote() for _ in range({num_cpus_per_node * expected_nodes})]))
|
|
print("end")
|
|
"""
|
|
|
|
try:
|
|
cluster.start()
|
|
ray.init("auto")
|
|
gcs_address = ray.get_runtime_context().gcs_address
|
|
|
|
def tasks_run():
|
|
tasks = list_tasks()
|
|
# Waiting til the driver in the run_string_as_driver_nonblocking is running
|
|
assert len(tasks) > 0
|
|
return True
|
|
|
|
run_string_as_driver_nonblocking(driver_script)
|
|
wait_for_condition(tasks_run)
|
|
|
|
reached_threshold = False
|
|
for _ in range(30):
|
|
# verify no pending task + with resource used.
|
|
status = get_cluster_status(gcs_address)
|
|
has_task_demand = len(status.resource_demands.ray_task_actor_demand) > 0
|
|
|
|
# Autoscaler can briefly launch extra workers while demand and
|
|
# in-flight instance views catch up; it is then idle-terminated.
|
|
# Check that we don't overscale (allow one transient extra node).
|
|
assert len(status.active_nodes) <= expected_nodes + 1
|
|
|
|
# Check there's no demand if we've reached the expected number of nodes
|
|
if reached_threshold:
|
|
assert not has_task_demand
|
|
|
|
# Load disappears in the next cycle after we've fully scaled up.
|
|
if len(status.active_nodes) >= expected_nodes:
|
|
reached_threshold = True
|
|
|
|
time.sleep(1)
|
|
|
|
assert reached_threshold
|
|
finally:
|
|
# TODO(rickyx): refactor into a fixture for autoscaling cluster.
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
# TODO(rickyx): We are NOT able to counter multi-node inconsistency yet. The problem is
|
|
# right now, when node A (head node) has an infeasible task,
|
|
# node B just finished running previous task.
|
|
# the actual cluster view will be:
|
|
# node A: 1 pending task (infeasible)
|
|
# node B: 0 pending task, CPU used = 0
|
|
#
|
|
# However, when node B's state is not updated on node A, the cluster view will be:
|
|
# node A: 1 pending task (infeasible)
|
|
# node B: 0 pending task, but **CPU used = 1**
|
|
#
|
|
@pytest.mark.parametrize("mode", (["single_node", "multi_node"]))
|
|
@pytest.mark.parametrize("autoscaler_v2", [False, True], ids=["v1", "v2"])
|
|
def test_scheduled_task_no_pending_demand(mode, autoscaler_v2):
|
|
|
|
# So that head node will need to dispatch tasks to worker node.
|
|
num_head_cpu = 0 if mode == "multi_node" else 1
|
|
|
|
cluster = AutoscalingCluster(
|
|
head_resources={"CPU": num_head_cpu},
|
|
worker_node_types={
|
|
"type-1": {
|
|
"resources": {"CPU": 1},
|
|
"node_config": {},
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
},
|
|
autoscaler_v2=autoscaler_v2,
|
|
)
|
|
|
|
driver_script = """
|
|
import time
|
|
import ray
|
|
@ray.remote(num_cpus=1)
|
|
def foo():
|
|
return True
|
|
|
|
ray.init("auto")
|
|
|
|
while True:
|
|
assert(ray.get(foo.remote()))
|
|
"""
|
|
|
|
try:
|
|
cluster.start()
|
|
ray.init("auto")
|
|
gcs_address = ray.get_runtime_context().gcs_address
|
|
|
|
run_string_as_driver_nonblocking(driver_script)
|
|
|
|
def tasks_run():
|
|
tasks = list_tasks()
|
|
|
|
# Waiting til the driver in the run_string_as_driver_nonblocking is running
|
|
assert len(tasks) > 0
|
|
|
|
return True
|
|
|
|
wait_for_condition(tasks_run)
|
|
|
|
for _ in range(30):
|
|
# verify no pending task + with resource used.
|
|
status = get_cluster_status(gcs_address)
|
|
has_task_demand = len(status.resource_demands.ray_task_actor_demand) > 0
|
|
has_task_usage = False
|
|
|
|
for usage in status.cluster_resource_usage:
|
|
if usage.resource_name == "CPU":
|
|
has_task_usage = usage.used > 0
|
|
print(status.cluster_resource_usage)
|
|
print(status.resource_demands.ray_task_actor_demand)
|
|
assert not (has_task_demand and has_task_usage), status
|
|
time.sleep(0.1)
|
|
finally:
|
|
# TODO(rickyx): refactor into a fixture for autoscaling cluster.
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
@pytest.mark.parametrize("autoscaler_v2", [False, True], ids=["v1", "v2"])
|
|
def test_placement_group_consistent(autoscaler_v2):
|
|
# Test that continuously creating and removing placement groups
|
|
# does not leak pending resource requests.
|
|
import time
|
|
|
|
cluster = AutoscalingCluster(
|
|
head_resources={"CPU": 0},
|
|
worker_node_types={
|
|
"type-1": {
|
|
"resources": {"CPU": 1},
|
|
"node_config": {},
|
|
"min_workers": 0,
|
|
"max_workers": 2,
|
|
},
|
|
},
|
|
autoscaler_v2=autoscaler_v2,
|
|
)
|
|
driver_script = """
|
|
|
|
import ray
|
|
import time
|
|
# Import placement group APIs.
|
|
from ray.util.placement_group import (
|
|
placement_group,
|
|
placement_group_table,
|
|
remove_placement_group,
|
|
)
|
|
|
|
ray.init("auto")
|
|
|
|
# Reserve all the CPUs of nodes, X= num of cpus, N = num of nodes
|
|
while True:
|
|
pg = placement_group([{"CPU": 1}])
|
|
ray.get(pg.ready())
|
|
time.sleep(0.5)
|
|
remove_placement_group(pg)
|
|
time.sleep(0.5)
|
|
"""
|
|
|
|
try:
|
|
cluster.start()
|
|
ray.init("auto")
|
|
gcs_address = ray.get_runtime_context().gcs_address
|
|
|
|
run_string_as_driver_nonblocking(driver_script)
|
|
|
|
def pg_created():
|
|
pgs = list_placement_groups()
|
|
assert len(pgs) > 0
|
|
|
|
return True
|
|
|
|
wait_for_condition(pg_created)
|
|
|
|
for _ in range(30):
|
|
# verify no pending request + resource used.
|
|
status = get_cluster_status(gcs_address)
|
|
has_pg_demand = len(status.resource_demands.placement_group_demand) > 0
|
|
has_pg_usage = False
|
|
for usage in status.cluster_resource_usage:
|
|
has_pg_usage = has_pg_usage or "bundle" in usage.resource_name
|
|
print(has_pg_demand, has_pg_usage)
|
|
assert not (has_pg_demand and has_pg_usage), status
|
|
time.sleep(0.1)
|
|
finally:
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
def test_autoscaler_v2_usage_report():
|
|
|
|
# Test that nodes become idle after placement group removal.
|
|
cluster = AutoscalingCluster(
|
|
head_resources={"CPU": 2},
|
|
worker_node_types={
|
|
"type-1": {
|
|
"resources": {"CPU": 2},
|
|
"node_config": {},
|
|
"min_workers": 0,
|
|
"max_workers": 2,
|
|
},
|
|
},
|
|
autoscaler_v2=True,
|
|
)
|
|
try:
|
|
cluster.start()
|
|
ray.init("auto")
|
|
gcs_client = GcsClient(ray.get_runtime_context().gcs_address)
|
|
|
|
def verify():
|
|
tags = get_extra_usage_tags_to_report(gcs_client)
|
|
print(tags)
|
|
assert tags[TagKey.Name(TagKey.AUTOSCALER_VERSION).lower()] == "v2", tags
|
|
return True
|
|
|
|
wait_for_condition(verify)
|
|
finally:
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
@pytest.mark.parametrize("autoscaler_v2", [False, True], ids=["v1", "v2"])
|
|
def test_placement_group_removal_idle_node(autoscaler_v2):
|
|
# Test that nodes become idle after placement group removal.
|
|
cluster = AutoscalingCluster(
|
|
head_resources={"CPU": 2},
|
|
worker_node_types={
|
|
"type-1": {
|
|
"resources": {"CPU": 2},
|
|
"node_config": {},
|
|
"min_workers": 0,
|
|
"max_workers": 2,
|
|
},
|
|
},
|
|
autoscaler_v2=autoscaler_v2,
|
|
)
|
|
try:
|
|
cluster.start()
|
|
ray.init("auto")
|
|
gcs_address = ray.get_runtime_context().gcs_address
|
|
|
|
# Schedule a pg on nodes
|
|
pg = placement_group([{"CPU": 2}] * 3, strategy="STRICT_SPREAD")
|
|
ray.get(pg.ready())
|
|
|
|
time.sleep(2)
|
|
remove_placement_group(pg)
|
|
|
|
from ray.autoscaler.v2.sdk import get_cluster_status
|
|
|
|
def verify():
|
|
cluster_state = get_cluster_status(gcs_address)
|
|
|
|
# Verify that nodes are idle.
|
|
assert len((cluster_state.idle_nodes)) == 3
|
|
for node in cluster_state.idle_nodes:
|
|
assert node.node_status == "IDLE"
|
|
assert node.resource_usage.idle_time_ms >= 1000
|
|
|
|
return True
|
|
|
|
wait_for_condition(verify)
|
|
finally:
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
@pytest.mark.parametrize("autoscaler_v2", [False, True], ids=["v1", "v2"])
|
|
def test_placement_group_reschedule_node_dead(autoscaler_v2):
|
|
# Test autoscaler reschedules placement group when node dies.
|
|
# Note that it should only provision nodes for the bundles that haven't been placed.
|
|
|
|
cluster = AutoscalingCluster(
|
|
head_resources={"CPU": 0},
|
|
worker_node_types={
|
|
"type-1": {
|
|
"resources": {"R1": 1},
|
|
"node_config": {},
|
|
"min_workers": 0,
|
|
"max_workers": 2,
|
|
},
|
|
"type-2": {
|
|
"resources": {"R2": 1},
|
|
"node_config": {},
|
|
"min_workers": 0,
|
|
"max_workers": 2,
|
|
},
|
|
"type-3": {
|
|
"resources": {"R3": 1},
|
|
"node_config": {},
|
|
"min_workers": 0,
|
|
"max_workers": 2,
|
|
},
|
|
},
|
|
autoscaler_v2=autoscaler_v2,
|
|
)
|
|
|
|
try:
|
|
cluster.start()
|
|
ray.init("auto")
|
|
gcs_address = ray.get_runtime_context().gcs_address
|
|
|
|
pg = placement_group([{"R1": 1}, {"R2": 1}, {"R3": 1}])
|
|
|
|
ray.get(pg.ready())
|
|
|
|
def verify_nodes(active, idle):
|
|
cluster_state = get_cluster_status(gcs_address)
|
|
assert len(cluster_state.active_nodes) == active
|
|
assert len(cluster_state.idle_nodes) == idle
|
|
return True
|
|
|
|
# 3 worker nodes, 1 head node (idle)
|
|
wait_for_condition(lambda: verify_nodes(3, 1))
|
|
|
|
def kill_node(node_id):
|
|
cmd = f"ps auxww | grep {node_id} | grep -v grep | awk '{{print $2}}'"
|
|
pids = subprocess.check_output(cmd, shell=True).decode("utf-8").strip()
|
|
print(f"Killing pids {pids}")
|
|
# kill the pids (handle multiple PIDs separated by newlines)
|
|
for pid in pids.split("\n"):
|
|
if pid:
|
|
cmd = f"kill -9 {pid}"
|
|
subprocess.run(cmd, shell=True)
|
|
|
|
# Kill a worker node with 'R1' in resources
|
|
for n in ray.nodes():
|
|
if "R1" in n["Resources"]:
|
|
node = n
|
|
break
|
|
|
|
# TODO(mimi): kill_raylet won't trigger reschedule in autoscaler v1
|
|
kill_node(node["NodeID"])
|
|
|
|
# Wait for the node to be removed
|
|
wait_for_condition(lambda: verify_nodes(2, 1), 30)
|
|
|
|
# Only provision nodes for unplaced bundles;
|
|
# avoid rescheduling the whole placement group.
|
|
wait_for_condition(lambda: verify_nodes(3, 1))
|
|
|
|
# Verify that the R1 node is recreated and has a different NodeID.
|
|
assert any(
|
|
[
|
|
"R1" in n["Resources"] and node["NodeID"] != n["NodeID"]
|
|
for n in ray.nodes()
|
|
]
|
|
), "R1 node is not recreated."
|
|
|
|
finally:
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
def test_object_store_memory_idle_node(shutdown_only):
|
|
|
|
ray.init()
|
|
|
|
obj = ray.put("hello")
|
|
gcs_address = ray.get_runtime_context().gcs_address
|
|
|
|
def verify():
|
|
state = get_cluster_status(gcs_address)
|
|
for node in state.active_nodes:
|
|
assert node.node_status == "RUNNING"
|
|
assert node.used_resources()["object_store_memory"] > 0
|
|
assert len(state.idle_nodes) == 0
|
|
return True
|
|
|
|
wait_for_condition(verify)
|
|
|
|
del obj
|
|
|
|
import time
|
|
|
|
time.sleep(1)
|
|
|
|
def verify():
|
|
state = get_cluster_status(gcs_address)
|
|
for node in state.idle_nodes:
|
|
assert node.node_status == "IDLE"
|
|
assert node.used_resources()["object_store_memory"] == 0
|
|
assert node.resource_usage.idle_time_ms >= 1000
|
|
assert len(state.active_nodes) == 0
|
|
return True
|
|
|
|
wait_for_condition(verify)
|
|
|
|
|
|
@pytest.mark.parametrize("autoscaler_v2", [False, True], ids=["v1", "v2"])
|
|
def test_non_corrupted_resources(autoscaler_v2):
|
|
"""
|
|
Test that when node's local gc happens due to object store pressure,
|
|
the message doesn't corrupt the resource view on the gcs.
|
|
See issue https://github.com/ray-project/ray/issues/39644
|
|
"""
|
|
num_worker_nodes = 5
|
|
cluster = AutoscalingCluster(
|
|
head_resources={"CPU": 2, "object_store_memory": 100 * 1024 * 1024},
|
|
worker_node_types={
|
|
"type-1": {
|
|
"resources": {"CPU": 2},
|
|
"node_config": {},
|
|
"min_workers": num_worker_nodes,
|
|
"max_workers": num_worker_nodes,
|
|
},
|
|
},
|
|
idle_timeout_minutes=999,
|
|
autoscaler_v2=autoscaler_v2,
|
|
)
|
|
|
|
driver_script = """
|
|
|
|
import ray
|
|
import time
|
|
|
|
ray.init("auto")
|
|
|
|
@ray.remote(num_cpus=1)
|
|
def foo():
|
|
ray.put(bytearray(1024*1024* 50))
|
|
|
|
|
|
while True:
|
|
ray.get([foo.remote() for _ in range(50)])
|
|
"""
|
|
|
|
try:
|
|
# This should trigger many COMMANDS messages from NodeManager.
|
|
cluster.start(
|
|
_system_config={
|
|
"debug_dump_period_milliseconds": 10,
|
|
"raylet_report_resources_period_milliseconds": 10000,
|
|
"global_gc_min_interval_s": 1,
|
|
"local_gc_interval_s": 1,
|
|
"plasma_store_usage_trigger_gc_threshold": 0.2,
|
|
"raylet_check_gc_period_milliseconds": 10,
|
|
},
|
|
)
|
|
ctx = ray.init("auto")
|
|
gcs_address = ctx.address_info["gcs_address"]
|
|
|
|
from ray.autoscaler.v2.sdk import get_cluster_status
|
|
|
|
def nodes_up():
|
|
cluster_state = get_cluster_status(gcs_address)
|
|
assert len(cluster_state.idle_nodes) == num_worker_nodes + 1
|
|
return True
|
|
|
|
wait_for_condition(nodes_up, timeout=20)
|
|
|
|
# Schedule tasks
|
|
run_string_as_driver_nonblocking(driver_script)
|
|
start = time.time()
|
|
|
|
# Check the cluster state for 10 seconds
|
|
while time.time() - start < 10:
|
|
cluster_state = get_cluster_status(gcs_address)
|
|
|
|
# Verify total cluster resources never change
|
|
assert (
|
|
len(cluster_state.idle_nodes) + len(cluster_state.active_nodes)
|
|
) == num_worker_nodes + 1
|
|
assert cluster_state.total_resources()["CPU"] == 2 * (num_worker_nodes + 1)
|
|
|
|
finally:
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
# Helper function to vaidate that a node's labels satisfy a `label_selector`.
|
|
def _verify_node_labels_for_selector(
|
|
node_labels: Dict[str, str], selector: Dict[str, str]
|
|
) -> bool:
|
|
for key, value in selector.items():
|
|
node_val = node_labels.get(key)
|
|
|
|
if "!in(" in value:
|
|
options_str = value.replace("!in(", "").replace(")", "")
|
|
options = {opt.strip() for opt in options_str.split(",")}
|
|
if node_val in options:
|
|
return False
|
|
elif "in(" in value:
|
|
options_str = value.replace("in(", "").replace(")", "")
|
|
options = {opt.strip() for opt in options_str.split(",")}
|
|
if node_val not in options:
|
|
return False
|
|
elif value.startswith("!"):
|
|
if node_val == value[1:]:
|
|
return False
|
|
else:
|
|
if node_val != value:
|
|
return False
|
|
# If all checks pass for all key-value pairs in the selector, return True.
|
|
return True
|
|
|
|
|
|
@pytest.mark.parametrize("autoscaler_v2", [True])
|
|
def test_task_scheduled_on_node_with_label_selector(autoscaler_v2):
|
|
cluster = AutoscalingCluster(
|
|
head_resources={"CPU": 0},
|
|
worker_node_types={
|
|
"node1": {
|
|
"resources": {"CPU": 1},
|
|
"node_config": {},
|
|
"labels": {"accelerator-type": "A100", "market-type": "spot"},
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
"node2": {
|
|
"resources": {"CPU": 1},
|
|
"node_config": {},
|
|
"labels": {
|
|
"region": "us-east1",
|
|
"accelerator-type": "TPU",
|
|
"market-type": "spot",
|
|
},
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
"node3": {
|
|
"resources": {"CPU": 1},
|
|
"node_config": {},
|
|
"labels": {"accelerator-type": "B200", "market-type": "spot"},
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
"node4": {
|
|
"resources": {"CPU": 1},
|
|
"node_config": {},
|
|
"labels": {"market-type": "on-demand", "accelerator-type": "TPU"},
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
},
|
|
idle_timeout_minutes=999,
|
|
autoscaler_v2=autoscaler_v2,
|
|
)
|
|
|
|
driver_script = """
|
|
import ray
|
|
import time
|
|
|
|
@ray.remote(num_cpus=1)
|
|
def labels_task():
|
|
time.sleep(20)
|
|
return True
|
|
|
|
ray.init("auto")
|
|
|
|
label_selectors = [
|
|
{"accelerator-type": "A100"},
|
|
{"region": "in(us-east1,me-central1)"},
|
|
{"accelerator-type": "!in(A100,TPU)"},
|
|
{"market-type": "!spot"},
|
|
]
|
|
|
|
results = [
|
|
labels_task.options(name=f"task_{i}", label_selector=sel).remote()
|
|
for i, sel in enumerate(label_selectors)
|
|
]
|
|
assert all(ray.get(results))
|
|
"""
|
|
|
|
try:
|
|
cluster.start()
|
|
ray.init("auto")
|
|
gcs_address = ray.get_runtime_context().gcs_address
|
|
expected_nodes = 4
|
|
|
|
def all_tasks_submitted():
|
|
return len(list_tasks()) == expected_nodes
|
|
|
|
proc = run_string_as_driver_nonblocking(driver_script)
|
|
wait_for_condition(all_tasks_submitted)
|
|
|
|
def all_nodes_launched():
|
|
status = get_cluster_status(gcs_address)
|
|
return len(status.active_nodes) == expected_nodes
|
|
|
|
wait_for_condition(all_nodes_launched, timeout=30)
|
|
proc.wait(timeout=30)
|
|
assert proc.returncode == 0, "The driver script failed."
|
|
|
|
# Validate Tasks are scheduled on nodes with required labels.
|
|
tasks_by_name = {
|
|
task.name: task for task in list_tasks(detail=True) if hasattr(task, "name")
|
|
}
|
|
nodes = {node["NodeID"]: node["Labels"] for node in ray.nodes()}
|
|
task_selectors = {
|
|
"task_0": {"accelerator-type": "A100"},
|
|
"task_1": {"region": "in(me-central1,us-east1)"},
|
|
"task_2": {"accelerator-type": "!in(A100,TPU)"},
|
|
"task_3": {"market-type": "!spot"},
|
|
}
|
|
|
|
for task_name, expected_selector in task_selectors.items():
|
|
assert (
|
|
task_name in tasks_by_name
|
|
), f"Task with name '{task_name}' was not found."
|
|
task = tasks_by_name[task_name]
|
|
|
|
# Verify actual label selector from the Task matches the expected.
|
|
actual_selector = task.get("label_selector")
|
|
assert (
|
|
actual_selector is not None
|
|
), f"Task '{task_name}' did not have a 'label_selector' field."
|
|
|
|
assert actual_selector == expected_selector, (
|
|
f"Task '{task_name}' has an incorrect label selector. "
|
|
f"Expected: {expected_selector}, Got: {actual_selector}"
|
|
)
|
|
|
|
# Verify Ray node created for Task.
|
|
node_id = task["node_id"]
|
|
assert (
|
|
node_id in nodes
|
|
), f"Node with ID '{node_id}' for task '{task_name}' was not found."
|
|
|
|
# Validate node labels satisfy `label_selector` for Task.
|
|
node_labels = nodes[node_id]
|
|
assert _verify_node_labels_for_selector(
|
|
node_labels, actual_selector
|
|
), f"Verification failed for task '{task_name}' on node '{node_id}'"
|
|
|
|
finally:
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
@pytest.mark.parametrize("autoscaler_v2", [True])
|
|
def test_actor_scheduled_on_node_with_label_selector(autoscaler_v2):
|
|
cluster = AutoscalingCluster(
|
|
head_resources={"CPU": 0},
|
|
worker_node_types={
|
|
"node1": {
|
|
"resources": {"CPU": 1},
|
|
"node_config": {},
|
|
"labels": {"accelerator-type": "A100", "market-type": "spot"},
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
"node2": {
|
|
"resources": {"CPU": 1},
|
|
"node_config": {},
|
|
"labels": {
|
|
"region": "us-east1",
|
|
"accelerator-type": "TPU",
|
|
"market-type": "spot",
|
|
},
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
"node3": {
|
|
"resources": {"CPU": 1},
|
|
"node_config": {},
|
|
"labels": {"accelerator-type": "B200", "market-type": "spot"},
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
"node4": {
|
|
"resources": {"CPU": 1},
|
|
"node_config": {},
|
|
"labels": {"market-type": "on-demand", "accelerator-type": "TPU"},
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
},
|
|
idle_timeout_minutes=999,
|
|
autoscaler_v2=autoscaler_v2,
|
|
)
|
|
|
|
driver_script = """
|
|
import ray
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def ready(self):
|
|
return True
|
|
|
|
ray.init("auto")
|
|
|
|
label_selectors = [
|
|
{"accelerator-type": "A100"},
|
|
{"region": "in(us-east1,me-central1)"},
|
|
{"accelerator-type": "!in(A100,TPU)"},
|
|
{"market-type": "!spot"},
|
|
]
|
|
|
|
actors = [
|
|
Actor.options(name=f"actor_{i}", label_selector=sel).remote()
|
|
for i, sel in enumerate(label_selectors)
|
|
]
|
|
|
|
ray.get([a.ready.remote() for a in actors])
|
|
"""
|
|
|
|
try:
|
|
cluster.start()
|
|
ray.init("auto")
|
|
gcs_address = ray.get_runtime_context().gcs_address
|
|
expected_nodes = 4
|
|
|
|
def all_actors_submitted():
|
|
return len(list_actors()) == expected_nodes
|
|
|
|
proc = run_string_as_driver_nonblocking(driver_script)
|
|
wait_for_condition(all_actors_submitted)
|
|
|
|
def all_actors_scheduled():
|
|
# Verify the nodes launched for the Actors are as expected.
|
|
status = get_cluster_status(gcs_address)
|
|
if len(status.active_nodes) != expected_nodes:
|
|
return False
|
|
|
|
active_node_types = {
|
|
node.ray_node_type_name for node in status.active_nodes
|
|
}
|
|
expected_node_types = {"node1", "node2", "node3", "node4"}
|
|
return active_node_types == expected_node_types
|
|
|
|
# All Actors with label selectors should be scheduled, scaling
|
|
# 4 nodes with the required labels.
|
|
wait_for_condition(all_actors_scheduled, timeout=30)
|
|
proc.wait(timeout=30)
|
|
assert proc.returncode == 0, "The driver script failed to submit actors."
|
|
|
|
# Finally, validate the Actors are scheduled on the node with matching labels.
|
|
actors_by_name = {
|
|
actor.name: actor
|
|
for actor in list_actors(detail=True)
|
|
if hasattr(actor, "name")
|
|
}
|
|
nodes = {node["NodeID"]: node["Labels"] for node in ray.nodes()}
|
|
actor_selectors = {
|
|
"actor_0": {"accelerator-type": "A100"},
|
|
"actor_1": {"region": "in(me-central1,us-east1)"},
|
|
"actor_2": {"accelerator-type": "!in(A100,TPU)"},
|
|
"actor_3": {"market-type": "!spot"},
|
|
}
|
|
|
|
for actor_name, expected_selector in actor_selectors.items():
|
|
assert (
|
|
actor_name in actors_by_name
|
|
), f"Actor with name '{actor_name}' was not found."
|
|
actor = actors_by_name[actor_name]
|
|
|
|
# Verify actual label selector from the Actor matches the expected.
|
|
actual_selector = actor.get("label_selector")
|
|
assert (
|
|
actual_selector is not None
|
|
), f"Actor '{actor_name}' did not have a 'label_selector' field."
|
|
|
|
assert actual_selector == expected_selector, (
|
|
f"Actor '{actor_name}' has an incorrect label selector. "
|
|
f"Expected: {expected_selector}, Got: {actual_selector}"
|
|
)
|
|
|
|
# Verify Ray node created for Actor.
|
|
node_id = actor["node_id"]
|
|
assert (
|
|
node_id in nodes
|
|
), f"Node with ID '{node_id}' for Actor '{actor_name}' was not found."
|
|
|
|
# Validate node labels satisfy `label_selector` for Actor.
|
|
node_labels = nodes[node_id]
|
|
assert _verify_node_labels_for_selector(
|
|
node_labels, actual_selector
|
|
), f"Verification failed for Actor '{actor_name}' on node '{node_id}'"
|
|
|
|
finally:
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
@pytest.mark.parametrize("autoscaler_v2", [True])
|
|
def test_pg_scheduled_on_node_with_bundle_label_selector(autoscaler_v2):
|
|
cluster = AutoscalingCluster(
|
|
head_resources={"CPU": 0},
|
|
worker_node_types={
|
|
"unlabelled_node": {
|
|
"resources": {"CPU": 1, "GPU": 1, "TPU": 1},
|
|
"node_config": {},
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
"not_matching_labels": {
|
|
"resources": {"CPU": 1},
|
|
"labels": {"unrelated": "labels"},
|
|
"node_config": {},
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
"a100_node": {
|
|
"resources": {"CPU": 1, "GPU": 1},
|
|
"node_config": {},
|
|
"labels": {"accelerator-type": "A100"},
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
"tpu_node": {
|
|
"resources": {"CPU": 1, "TPU": 1},
|
|
"node_config": {},
|
|
"labels": {"accelerator-type": "TPU_V6E"},
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
},
|
|
idle_timeout_minutes=999,
|
|
autoscaler_v2=autoscaler_v2,
|
|
)
|
|
|
|
try:
|
|
cluster.start()
|
|
ray.init("auto")
|
|
gcs_address = ray.get_runtime_context().gcs_address
|
|
# We expect one GPU and one TPU node to scale.
|
|
expected_nodes = 2
|
|
|
|
# Define a placement group where each bundle should scale a node of a different type.
|
|
pg = placement_group(
|
|
name="label_selector_pg",
|
|
bundles=[
|
|
{"CPU": 1},
|
|
{"CPU": 1},
|
|
],
|
|
bundle_label_selector=[
|
|
{"accelerator-type": "A100"}, # a100_node
|
|
{"accelerator-type": "TPU_V6E"}, # tpu_node
|
|
],
|
|
strategy="SPREAD",
|
|
)
|
|
|
|
# Wait for the placement group to be ready.
|
|
ray.get(pg.ready())
|
|
|
|
# Validate the number and types of the auto-scaled nodes are as expected.
|
|
# Add a wait here to avoid flaky test behavior.
|
|
def check_nodes_active():
|
|
status = get_cluster_status(gcs_address)
|
|
return len(status.active_nodes) == expected_nodes
|
|
|
|
try:
|
|
wait_for_condition(check_nodes_active, timeout=30, retry_interval_ms=500)
|
|
except Exception as e:
|
|
latest_status = get_cluster_status(gcs_address)
|
|
raise AssertionError(
|
|
f"Timed out waiting for {expected_nodes} active nodes. "
|
|
f"Got: {len(latest_status.active_nodes)}. "
|
|
f"Full status: {latest_status}"
|
|
) from e
|
|
|
|
status = get_cluster_status(gcs_address)
|
|
actual_node_types = {node.ray_node_type_name for node in status.active_nodes}
|
|
expected_node_types = {"a100_node", "tpu_node"}
|
|
assert actual_node_types == expected_node_types
|
|
|
|
# Validate the placement group is scheduled to nodes with the required labels.
|
|
pgs = list_placement_groups(detail=True)
|
|
assert len(pgs) == 1
|
|
pg_state = pgs[0]
|
|
bundles_list = pg_state.bundles
|
|
assert (
|
|
bundles_list is not None
|
|
), "PlacementGroupState did not have a 'bundles' field."
|
|
|
|
actual_bundle_selectors = []
|
|
for bundle in bundles_list:
|
|
actual_bundle_selectors.append(bundle["label_selector"])
|
|
|
|
expected_bundle_selectors = [
|
|
{"accelerator-type": "A100"},
|
|
{"accelerator-type": "TPU_V6E"},
|
|
]
|
|
assert actual_bundle_selectors == expected_bundle_selectors, (
|
|
f"Placement group has incorrect bundle selectors. "
|
|
f"Expected: {expected_bundle_selectors}, Got: {actual_bundle_selectors}"
|
|
)
|
|
|
|
nodes = {node["NodeID"]: node["Labels"] for node in ray.nodes()}
|
|
for bundle_index, bundle in enumerate(bundles_list):
|
|
# Verify bundle placed on expected node.
|
|
bundle_node_id = bundle.get("node_id")
|
|
assert (
|
|
bundle_node_id in nodes
|
|
), f"Node with ID '{bundle_node_id}' for bundle {bundle_index} was not found."
|
|
|
|
# Verify node's labels satisfy the bundle's label_selector.
|
|
bundle_selector = actual_bundle_selectors[bundle_index]
|
|
node_labels = nodes[bundle_node_id]
|
|
assert _verify_node_labels_for_selector(node_labels, bundle_selector)
|
|
|
|
finally:
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
def test_priority_selection_e2e():
|
|
cluster = AutoscalingCluster(
|
|
head_resources={"CPU": 0},
|
|
worker_node_types={
|
|
"high-priority": {
|
|
"resources": {"CPU": 1},
|
|
"node_config": {},
|
|
"priority": 10,
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
"low-priority": {
|
|
"resources": {"CPU": 1},
|
|
"node_config": {},
|
|
"priority": 0,
|
|
"min_workers": 0,
|
|
"max_workers": 1,
|
|
},
|
|
},
|
|
autoscaler_v2=True,
|
|
)
|
|
|
|
try:
|
|
cluster.start()
|
|
ray.init("auto")
|
|
gcs_address = ray.get_runtime_context().gcs_address
|
|
|
|
@ray.remote(num_cpus=1)
|
|
def foo():
|
|
import time
|
|
|
|
time.sleep(5)
|
|
return True
|
|
|
|
# Submit a task
|
|
foo.remote()
|
|
|
|
def high_priority_node_launched():
|
|
status = get_cluster_status(gcs_address)
|
|
active_node_types = {
|
|
node.ray_node_type_name for node in status.active_nodes
|
|
}
|
|
assert "low-priority" not in active_node_types
|
|
return "high-priority" in active_node_types
|
|
|
|
# Wait for the high priority node to be launched
|
|
wait_for_condition(high_priority_node_launched, timeout=30)
|
|
|
|
finally:
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
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__]))
|