wehub-resource-sync
894 lines
29 KiB
Python
894 lines
29 KiB
Python
import os
|
|
import sys
|
|
import time
|
|
from typing import Dict, Optional
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
import ray._private.gcs_utils as gcs_utils
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray._private.test_utils import (
|
|
make_global_state_accessor,
|
|
)
|
|
from ray._raylet import GcsClient
|
|
from ray.core.generated import autoscaler_pb2
|
|
from ray.util.state import list_actors
|
|
|
|
|
|
def test_replenish_resources(ray_start_regular):
|
|
cluster_resources = ray.cluster_resources()
|
|
available_resources = ray.available_resources()
|
|
assert cluster_resources == available_resources
|
|
|
|
@ray.remote
|
|
def cpu_task():
|
|
pass
|
|
|
|
ray.get(cpu_task.remote())
|
|
|
|
wait_for_condition(lambda: ray.available_resources() == cluster_resources)
|
|
|
|
|
|
def test_uses_resources(ray_start_regular):
|
|
cluster_resources = ray.cluster_resources()
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
pass
|
|
|
|
actor = Actor.remote()
|
|
ray.get(actor.__ray_ready__.remote())
|
|
|
|
wait_for_condition(
|
|
lambda: ray.available_resources().get("CPU", 0)
|
|
== cluster_resources.get("CPU", 0) - 1
|
|
)
|
|
|
|
|
|
def test_available_resources_per_node(ray_start_cluster_head):
|
|
cluster = ray_start_cluster_head
|
|
|
|
@ray.remote
|
|
def get_node_id():
|
|
return ray.get_runtime_context().get_node_id()
|
|
|
|
head_node_id = ray.get(get_node_id.remote())
|
|
|
|
worker_node = cluster.add_node(num_cpus=3, resources={"worker": 1})
|
|
|
|
@ray.remote(num_cpus=1, resources={"worker": 1})
|
|
class Actor:
|
|
def ping(self):
|
|
return ray.get_runtime_context().get_node_id()
|
|
|
|
actor = Actor.remote()
|
|
worker_node_id = ray.get(actor.ping.remote())
|
|
|
|
def available_resources_per_node_check1():
|
|
available_resources_per_node = ray._private.state.available_resources_per_node()
|
|
assert len(available_resources_per_node) == 2
|
|
assert available_resources_per_node[head_node_id]["CPU"] == 1
|
|
assert available_resources_per_node[worker_node_id]["CPU"] == 2
|
|
assert available_resources_per_node[worker_node_id].get("worker", 0) == 0
|
|
return True
|
|
|
|
wait_for_condition(available_resources_per_node_check1)
|
|
|
|
cluster.remove_node(worker_node)
|
|
cluster.wait_for_nodes()
|
|
|
|
def available_resources_per_node_check2():
|
|
# Make sure worker node is not returned
|
|
available_resources_per_node = ray._private.state.available_resources_per_node()
|
|
assert len(available_resources_per_node) == 1
|
|
assert available_resources_per_node[head_node_id]["CPU"] == 1
|
|
return True
|
|
|
|
wait_for_condition(available_resources_per_node_check2)
|
|
|
|
|
|
def test_total_resources_per_node(ray_start_cluster_head):
|
|
cluster = ray_start_cluster_head
|
|
|
|
@ray.remote
|
|
def get_node_id():
|
|
return ray.get_runtime_context().get_node_id()
|
|
|
|
head_node_id = ray.get(get_node_id.remote())
|
|
|
|
worker_node = cluster.add_node(num_cpus=3, resources={"worker": 1})
|
|
|
|
@ray.remote(num_cpus=1, resources={"worker": 1})
|
|
class Actor:
|
|
def ping(self):
|
|
return ray.get_runtime_context().get_node_id()
|
|
|
|
actor = Actor.remote()
|
|
worker_node_id = ray.get(actor.ping.remote())
|
|
|
|
def total_resources_per_node_check1():
|
|
total_resources_per_node = ray._private.state.total_resources_per_node()
|
|
assert len(total_resources_per_node) == 2
|
|
assert total_resources_per_node[head_node_id]["CPU"] == 1
|
|
assert total_resources_per_node[worker_node_id]["CPU"] == 3
|
|
assert total_resources_per_node[worker_node_id].get("worker", 0) == 1
|
|
return True
|
|
|
|
wait_for_condition(total_resources_per_node_check1)
|
|
|
|
cluster.remove_node(worker_node)
|
|
cluster.wait_for_nodes()
|
|
|
|
def total_resources_per_node_check2():
|
|
# Make sure worker node is not returned
|
|
total_resources_per_node = ray._private.state.total_resources_per_node()
|
|
assert len(total_resources_per_node) == 1
|
|
assert total_resources_per_node[head_node_id]["CPU"] == 1
|
|
return True
|
|
|
|
wait_for_condition(total_resources_per_node_check2)
|
|
|
|
|
|
def test_add_remove_cluster_resources(ray_start_cluster_head):
|
|
"""Tests that Global State API is consistent with actual cluster."""
|
|
cluster = ray_start_cluster_head
|
|
assert ray.cluster_resources()["CPU"] == 1
|
|
nodes = []
|
|
nodes += [cluster.add_node(num_cpus=1)]
|
|
cluster.wait_for_nodes()
|
|
assert ray.cluster_resources()["CPU"] == 2
|
|
|
|
cluster.remove_node(nodes.pop())
|
|
cluster.wait_for_nodes()
|
|
assert ray.cluster_resources()["CPU"] == 1
|
|
|
|
for i in range(5):
|
|
nodes += [cluster.add_node(num_cpus=1)]
|
|
cluster.wait_for_nodes()
|
|
assert ray.cluster_resources()["CPU"] == 6
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_start_regular",
|
|
[{"include_dashboard": True}],
|
|
indirect=True,
|
|
)
|
|
def test_global_state_actor_table(ray_start_regular):
|
|
@ray.remote
|
|
class Actor:
|
|
def ready(self):
|
|
return os.getpid()
|
|
|
|
# actor table should be empty at first
|
|
assert len(list_actors()) == 0
|
|
|
|
a = Actor.remote()
|
|
pid = ray.get(a.ready.remote())
|
|
assert len(list_actors()) == 1
|
|
assert list_actors()[0].pid == pid
|
|
|
|
# actor table should contain only this entry
|
|
# even when the actor goes out of scope
|
|
del a
|
|
|
|
for _ in range(10):
|
|
if list_actors()[0].state == "DEAD":
|
|
break
|
|
else:
|
|
time.sleep(0.5)
|
|
assert list_actors()[0].state == "DEAD"
|
|
|
|
|
|
def test_global_state_worker_table(ray_start_regular):
|
|
def worker_initialized():
|
|
# Get worker table from gcs.
|
|
workers_data = ray._private.state.workers()
|
|
return len(workers_data) == 1
|
|
|
|
wait_for_condition(worker_initialized)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_start_regular",
|
|
[{"include_dashboard": True}],
|
|
indirect=True,
|
|
)
|
|
def test_global_state_actor_entry(ray_start_regular):
|
|
@ray.remote
|
|
class Actor:
|
|
def ready(self):
|
|
pass
|
|
|
|
# actor table should be empty at first
|
|
assert len(list_actors()) == 0
|
|
|
|
a = Actor.remote()
|
|
b = Actor.remote()
|
|
ray.get(a.ready.remote())
|
|
ray.get(b.ready.remote())
|
|
assert len(list_actors()) == 2
|
|
a_actor_id = a._actor_id.hex()
|
|
b_actor_id = b._actor_id.hex()
|
|
assert ray.util.state.get_actor(id=a_actor_id).actor_id == a_actor_id
|
|
assert ray.util.state.get_actor(id=a_actor_id).state == "ALIVE"
|
|
assert ray.util.state.get_actor(id=b_actor_id).actor_id == b_actor_id
|
|
assert ray.util.state.get_actor(id=b_actor_id).state == "ALIVE"
|
|
|
|
|
|
def test_node_name_cluster(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(node_name="head_node", include_dashboard=False)
|
|
head_context = ray.init(address=cluster.address, include_dashboard=False)
|
|
cluster.add_node(node_name="worker_node", include_dashboard=False)
|
|
cluster.wait_for_nodes()
|
|
|
|
global_state_accessor = make_global_state_accessor(head_context)
|
|
node_table = global_state_accessor.get_node_table()
|
|
assert len(node_table) == 2
|
|
for node in node_table:
|
|
if node["NodeID"] == head_context.address_info["node_id"]:
|
|
assert node["NodeName"] == "head_node"
|
|
else:
|
|
assert node["NodeName"] == "worker_node"
|
|
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
def test_node_name_init():
|
|
# Test ray.init with _node_name directly
|
|
new_head_context = ray.init(_node_name="new_head_node", include_dashboard=False)
|
|
|
|
global_state_accessor = make_global_state_accessor(new_head_context)
|
|
node = global_state_accessor.get_node_table()[0]
|
|
assert node["NodeName"] == "new_head_node"
|
|
ray.shutdown()
|
|
|
|
|
|
def test_no_node_name():
|
|
# Test that starting ray with no node name will result in a node_name=ip_address
|
|
new_head_context = ray.init(include_dashboard=False)
|
|
global_state_accessor = make_global_state_accessor(new_head_context)
|
|
node = global_state_accessor.get_node_table()[0]
|
|
assert node["NodeName"] == ray.util.get_node_ip_address()
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.mark.parametrize("max_shapes", [0, 2, -1])
|
|
def test_load_report(shutdown_only, max_shapes):
|
|
resource1 = "A"
|
|
resource2 = "B"
|
|
cluster = ray.init(
|
|
num_cpus=1,
|
|
resources={resource1: 1},
|
|
_system_config={
|
|
"max_resource_shapes_per_load_report": max_shapes,
|
|
},
|
|
)
|
|
|
|
global_state_accessor = make_global_state_accessor(cluster)
|
|
|
|
@ray.remote
|
|
def sleep():
|
|
time.sleep(1000)
|
|
|
|
sleep.remote()
|
|
for _ in range(3):
|
|
sleep.remote()
|
|
sleep.options(resources={resource1: 1}).remote()
|
|
sleep.options(resources={resource2: 1}).remote()
|
|
|
|
class Checker:
|
|
def __init__(self):
|
|
self.report = None
|
|
|
|
def check_load_report(self):
|
|
message = global_state_accessor.get_all_resource_usage()
|
|
if message is None:
|
|
return False
|
|
|
|
resource_usage = gcs_utils.ResourceUsageBatchData.FromString(message)
|
|
self.report = resource_usage.resource_load_by_shape.resource_demands
|
|
if max_shapes == 0:
|
|
return True
|
|
elif max_shapes == 2:
|
|
return len(self.report) >= 2
|
|
else:
|
|
return len(self.report) >= 3
|
|
|
|
# Wait for load information to arrive.
|
|
checker = Checker()
|
|
wait_for_condition(checker.check_load_report)
|
|
|
|
# Check that we respect the max shapes limit.
|
|
if max_shapes != -1:
|
|
assert len(checker.report) <= max_shapes
|
|
|
|
print(checker.report)
|
|
|
|
if max_shapes > 0:
|
|
# Check that we differentiate between infeasible and ready tasks.
|
|
for demand in checker.report:
|
|
if resource2 in demand.shape:
|
|
assert demand.num_infeasible_requests_queued > 0
|
|
assert demand.num_ready_requests_queued == 0
|
|
else:
|
|
assert demand.num_ready_requests_queued > 0
|
|
assert demand.num_infeasible_requests_queued == 0
|
|
|
|
|
|
def test_placement_group_load_report(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
# Add a head node that doesn't have gpu resource.
|
|
cluster.add_node(num_cpus=4)
|
|
|
|
global_state_accessor = make_global_state_accessor(
|
|
ray.init(address=cluster.address)
|
|
)
|
|
|
|
class PgLoadChecker:
|
|
def nothing_is_ready(self):
|
|
resource_usage = self._read_resource_usage()
|
|
if not resource_usage:
|
|
return False
|
|
if resource_usage.HasField("placement_group_load"):
|
|
pg_load = resource_usage.placement_group_load
|
|
return len(pg_load.placement_group_data) == 2
|
|
return False
|
|
|
|
def only_first_one_ready(self):
|
|
resource_usage = self._read_resource_usage()
|
|
if not resource_usage:
|
|
return False
|
|
if resource_usage.HasField("placement_group_load"):
|
|
pg_load = resource_usage.placement_group_load
|
|
return len(pg_load.placement_group_data) == 1
|
|
return False
|
|
|
|
def two_infeasible_pg(self):
|
|
resource_usage = self._read_resource_usage()
|
|
if not resource_usage:
|
|
return False
|
|
if resource_usage.HasField("placement_group_load"):
|
|
pg_load = resource_usage.placement_group_load
|
|
return len(pg_load.placement_group_data) == 2
|
|
return False
|
|
|
|
def _read_resource_usage(self):
|
|
message = global_state_accessor.get_all_resource_usage()
|
|
if message is None:
|
|
return False
|
|
|
|
resource_usage = gcs_utils.ResourceUsageBatchData.FromString(message)
|
|
return resource_usage
|
|
|
|
checker = PgLoadChecker()
|
|
|
|
# Create 2 placement groups that are infeasible.
|
|
pg_feasible = ray.util.placement_group([{"A": 1}])
|
|
pg_infeasible = ray.util.placement_group([{"B": 1}])
|
|
_, unready = ray.wait([pg_feasible.ready(), pg_infeasible.ready()], timeout=0)
|
|
assert len(unready) == 2
|
|
wait_for_condition(checker.nothing_is_ready)
|
|
|
|
# Add a node that makes pg feasible. Make sure load include this change.
|
|
cluster.add_node(resources={"A": 1})
|
|
ray.get(pg_feasible.ready())
|
|
wait_for_condition(checker.only_first_one_ready)
|
|
# Create one more infeasible pg and make sure load is properly updated.
|
|
pg_infeasible_second = ray.util.placement_group([{"C": 1}])
|
|
_, unready = ray.wait([pg_infeasible_second.ready()], timeout=0)
|
|
assert len(unready) == 1
|
|
wait_for_condition(checker.two_infeasible_pg)
|
|
|
|
|
|
def test_backlog_report(shutdown_only):
|
|
cluster = ray.init(
|
|
num_cpus=1,
|
|
_system_config={
|
|
"max_pending_lease_requests_per_scheduling_category": 1,
|
|
"report_worker_backlog_interval_ms": 100,
|
|
},
|
|
)
|
|
|
|
global_state_accessor = make_global_state_accessor(cluster)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
def foo(x):
|
|
print(".")
|
|
time.sleep(x)
|
|
return None
|
|
|
|
def backlog_size_set():
|
|
message = global_state_accessor.get_all_resource_usage()
|
|
if message is None:
|
|
return False
|
|
|
|
resource_usage = gcs_utils.ResourceUsageBatchData.FromString(message)
|
|
aggregate_resource_load = resource_usage.resource_load_by_shape.resource_demands
|
|
if len(aggregate_resource_load) == 1:
|
|
backlog_size = aggregate_resource_load[0].backlog_size
|
|
print(backlog_size)
|
|
# Ideally we'd want to assert backlog_size == 8, but guaranteeing
|
|
# the order the order that submissions will occur is too
|
|
# hard/flaky.
|
|
return backlog_size > 0
|
|
return False
|
|
|
|
# We want this first task to finish
|
|
refs = [foo.remote(0.5)]
|
|
# These tasks should all start _before_ the first one finishes.
|
|
refs.extend([foo.remote(1000) for _ in range(9)])
|
|
# Now there's 1 request running, 1 queued in the raylet, and 8 queued in
|
|
# the worker backlog.
|
|
|
|
ray.get(refs[0])
|
|
# First request finishes, second request is now running, third lease
|
|
# request is sent to the raylet with backlog=7
|
|
|
|
wait_for_condition(backlog_size_set, timeout=2)
|
|
|
|
|
|
def test_default_load_reports(shutdown_only):
|
|
"""Despite the fact that default actors release their cpu after being
|
|
placed, they should still require 1 CPU for laod reporting purposes.
|
|
https://github.com/ray-project/ray/issues/26806
|
|
"""
|
|
cluster = ray.init(
|
|
num_cpus=0,
|
|
)
|
|
|
|
global_state_accessor = make_global_state_accessor(cluster)
|
|
|
|
@ray.remote
|
|
def foo():
|
|
return None
|
|
|
|
@ray.remote
|
|
class Foo:
|
|
pass
|
|
|
|
def actor_and_task_queued_together():
|
|
message = global_state_accessor.get_all_resource_usage()
|
|
if message is None:
|
|
return False
|
|
|
|
resource_usage = gcs_utils.ResourceUsageBatchData.FromString(message)
|
|
aggregate_resource_load = resource_usage.resource_load_by_shape.resource_demands
|
|
print(f"Num shapes {len(aggregate_resource_load)}")
|
|
if len(aggregate_resource_load) == 1:
|
|
num_infeasible = aggregate_resource_load[0].num_infeasible_requests_queued
|
|
print(f"num in shape {num_infeasible}")
|
|
# Ideally we'd want to assert backlog_size == 8, but guaranteeing
|
|
# the order the order that submissions will occur is too
|
|
# hard/flaky.
|
|
return num_infeasible == 2
|
|
return False
|
|
|
|
# Assign to variables to keep the ref counter happy.
|
|
handle = Foo.remote()
|
|
ref = foo.remote()
|
|
|
|
wait_for_condition(actor_and_task_queued_together, timeout=2)
|
|
|
|
# Do something with the variables so lint is happy.
|
|
del handle
|
|
del ref
|
|
|
|
|
|
def test_heartbeat_ip(shutdown_only):
|
|
cluster = ray.init(num_cpus=1)
|
|
global_state_accessor = make_global_state_accessor(cluster)
|
|
self_ip = ray.util.get_node_ip_address()
|
|
|
|
def self_ip_is_set():
|
|
message = global_state_accessor.get_all_resource_usage()
|
|
if message is None:
|
|
return False
|
|
|
|
resource_usage = gcs_utils.ResourceUsageBatchData.FromString(message)
|
|
resources_data = resource_usage.batch[0]
|
|
return resources_data.node_manager_address == self_ip
|
|
|
|
wait_for_condition(self_ip_is_set, timeout=2)
|
|
|
|
|
|
def test_next_job_id(ray_start_regular):
|
|
job_id_1 = ray._private.state.next_job_id()
|
|
job_id_2 = ray._private.state.next_job_id()
|
|
assert job_id_1.int() + 1 == job_id_2.int()
|
|
|
|
|
|
def test_get_cluster_config(shutdown_only):
|
|
ray.init(num_cpus=1)
|
|
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
|
|
|
|
cluster_config = ray._private.state.state.get_cluster_config()
|
|
assert cluster_config is None
|
|
|
|
cluster_config = autoscaler_pb2.ClusterConfig()
|
|
cluster_config.max_resources["CPU"] = 100
|
|
node_group_config = autoscaler_pb2.NodeGroupConfig()
|
|
node_group_config.name = "m5.large"
|
|
node_group_config.resources["CPU"] = 5
|
|
node_group_config.max_count = -1
|
|
cluster_config.node_group_configs.append(node_group_config)
|
|
gcs_client.report_cluster_config(cluster_config.SerializeToString())
|
|
assert ray._private.state.state.get_cluster_config() == cluster_config
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"description, cluster_config, num_cpu",
|
|
[
|
|
(
|
|
"should return 0 since empty config is provided",
|
|
autoscaler_pb2.ClusterConfig(),
|
|
0,
|
|
),
|
|
(
|
|
"should return 0 since no node_group_config is provided",
|
|
autoscaler_pb2.ClusterConfig(
|
|
max_resources={"CPU": 100},
|
|
),
|
|
0,
|
|
),
|
|
(
|
|
"should return 0 since no CPU is provided under node_group_configs",
|
|
autoscaler_pb2.ClusterConfig(
|
|
max_resources={"CPU": 100},
|
|
node_group_configs=[autoscaler_pb2.NodeGroupConfig(name="m5.large")],
|
|
),
|
|
0,
|
|
),
|
|
(
|
|
"should return None since 0 instance is provided under node_group_configs",
|
|
autoscaler_pb2.ClusterConfig(
|
|
max_resources={"CPU": 100},
|
|
node_group_configs=[
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
resources={"CPU": 50},
|
|
name="m5.large",
|
|
max_count=0,
|
|
)
|
|
],
|
|
),
|
|
0,
|
|
),
|
|
(
|
|
"should return max since max_count=-1 under node_group_configs",
|
|
autoscaler_pb2.ClusterConfig(
|
|
max_resources={"CPU": 100},
|
|
node_group_configs=[
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
resources={"CPU": 50},
|
|
name="m5.large",
|
|
max_count=-1,
|
|
)
|
|
],
|
|
),
|
|
sys.maxsize,
|
|
),
|
|
(
|
|
"should return the total under node_group_configs since it is less than max_resources",
|
|
autoscaler_pb2.ClusterConfig(
|
|
max_resources={"CPU": 100},
|
|
node_group_configs=[
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
resources={"CPU": 50},
|
|
name="m5.large",
|
|
max_count=1,
|
|
)
|
|
],
|
|
),
|
|
50,
|
|
),
|
|
(
|
|
"should return the total under max_resources since it is less than node_group_configs total",
|
|
autoscaler_pb2.ClusterConfig(
|
|
max_resources={"CPU": 30},
|
|
node_group_configs=[
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
resources={"CPU": 50},
|
|
name="m5.large",
|
|
max_count=1,
|
|
)
|
|
],
|
|
),
|
|
30,
|
|
),
|
|
(
|
|
"should return the total under node_group_configs - no max_resources",
|
|
autoscaler_pb2.ClusterConfig(
|
|
node_group_configs=[
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
resources={"CPU": 50},
|
|
name="m5.large",
|
|
max_count=1,
|
|
)
|
|
],
|
|
),
|
|
50,
|
|
),
|
|
(
|
|
"should return the total under node_group_configs - multiple node_group_config",
|
|
autoscaler_pb2.ClusterConfig(
|
|
node_group_configs=[
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
resources={"CPU": 50},
|
|
name="m5.large",
|
|
max_count=1,
|
|
),
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
resources={"CPU": 10},
|
|
name="m5.small",
|
|
max_count=4,
|
|
),
|
|
],
|
|
),
|
|
90,
|
|
),
|
|
],
|
|
)
|
|
def test_get_max_cpus_from_cluster_config(
|
|
shutdown_only,
|
|
description: str,
|
|
cluster_config: autoscaler_pb2.ClusterConfig,
|
|
num_cpu: Optional[int],
|
|
):
|
|
ray.init(num_cpus=1)
|
|
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
|
|
|
|
gcs_client.report_cluster_config(cluster_config.SerializeToString())
|
|
max_resources = ray._private.state.state.get_max_resources_from_cluster_config()
|
|
num_cpu_from_max_resources = max_resources.get("CPU", 0) if max_resources else 0
|
|
assert num_cpu_from_max_resources == num_cpu, description
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"description, cluster_config, expected_resources",
|
|
[
|
|
(
|
|
"should return CPU/GPU/TPU as None since empty config is provided",
|
|
autoscaler_pb2.ClusterConfig(),
|
|
None,
|
|
),
|
|
(
|
|
"should return CPU/GPU/TPU as None since no node_group_config is provided",
|
|
autoscaler_pb2.ClusterConfig(
|
|
max_resources={"CPU": 100, "memory": 1000},
|
|
),
|
|
None,
|
|
),
|
|
(
|
|
"should return CPU/GPU/TPU plus resources from node_group_configs",
|
|
autoscaler_pb2.ClusterConfig(
|
|
node_group_configs=[
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
name="m5.large",
|
|
resources={"CPU": 50, "memory": 500},
|
|
max_count=1,
|
|
)
|
|
],
|
|
),
|
|
{"CPU": 50, "memory": 500},
|
|
),
|
|
(
|
|
"should return resources from both node_group_configs and max_resources",
|
|
autoscaler_pb2.ClusterConfig(
|
|
max_resources={"GPU": 8},
|
|
node_group_configs=[
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
name="m5.large",
|
|
resources={"CPU": 50, "memory": 500},
|
|
max_count=1,
|
|
)
|
|
],
|
|
),
|
|
{
|
|
"CPU": 50,
|
|
"memory": 500,
|
|
}, # GPU and TPU are None because not in node_group_configs
|
|
),
|
|
(
|
|
"should return limited by max_resources when node_group total exceeds it",
|
|
autoscaler_pb2.ClusterConfig(
|
|
max_resources={"CPU": 30, "memory": 200},
|
|
node_group_configs=[
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
name="m5.large",
|
|
resources={"CPU": 50, "memory": 500},
|
|
max_count=1,
|
|
)
|
|
],
|
|
),
|
|
{"CPU": 30, "memory": 200},
|
|
),
|
|
(
|
|
"should return sys.maxsize when max_count=-1",
|
|
autoscaler_pb2.ClusterConfig(
|
|
node_group_configs=[
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
name="m5.large",
|
|
resources={"CPU": 50, "custom_resource": 10},
|
|
max_count=-1,
|
|
)
|
|
],
|
|
),
|
|
{
|
|
"CPU": sys.maxsize,
|
|
"custom_resource": sys.maxsize,
|
|
},
|
|
),
|
|
(
|
|
"should sum across multiple node_group_configs",
|
|
autoscaler_pb2.ClusterConfig(
|
|
node_group_configs=[
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
name="m5.large",
|
|
resources={"CPU": 50, "memory": 500},
|
|
max_count=1,
|
|
),
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
name="m5.small",
|
|
resources={"CPU": 10, "GPU": 1},
|
|
max_count=4,
|
|
),
|
|
],
|
|
),
|
|
{
|
|
"CPU": 90,
|
|
"GPU": 4,
|
|
"memory": 500,
|
|
}, # 50 + (10*4), 500 + 0
|
|
),
|
|
(
|
|
"should return 0 for resources with 0 count or 0 resources",
|
|
autoscaler_pb2.ClusterConfig(
|
|
node_group_configs=[
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
name="m5.large",
|
|
resources={"CPU": 50, "memory": 0},
|
|
max_count=0, # This makes all resources None
|
|
),
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
name="m5.small",
|
|
resources={"GPU": 1},
|
|
max_count=2,
|
|
),
|
|
],
|
|
),
|
|
{
|
|
"CPU": 0,
|
|
"GPU": 2,
|
|
"memory": 0,
|
|
}, # CPU is None due to max_count=0, GPU has valid count
|
|
),
|
|
(
|
|
"should discover all resource types including custom ones",
|
|
autoscaler_pb2.ClusterConfig(
|
|
max_resources={"TPU": 16, "special_resource": 100},
|
|
node_group_configs=[
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
name="gpu-node",
|
|
resources={
|
|
"CPU": 32,
|
|
"GPU": 8,
|
|
"memory": 1000,
|
|
"custom_accelerator": 4,
|
|
},
|
|
max_count=2,
|
|
),
|
|
autoscaler_pb2.NodeGroupConfig(
|
|
name="cpu-node",
|
|
resources={"CPU": 96, "memory": 2000, "disk": 500},
|
|
max_count=1,
|
|
),
|
|
],
|
|
),
|
|
{
|
|
"CPU": 160, # (32*2) + (96*1)
|
|
"GPU": 16, # (8*2) + 0
|
|
"memory": 4000, # (1000*2) + (2000*1)
|
|
"custom_accelerator": 8, # (4*2) + 0
|
|
"disk": 500, # 0 + (500*1)
|
|
},
|
|
),
|
|
],
|
|
)
|
|
def test_get_max_resources_from_cluster_config(
|
|
shutdown_only,
|
|
description: str,
|
|
cluster_config: autoscaler_pb2.ClusterConfig,
|
|
expected_resources: Dict[str, Optional[int]],
|
|
):
|
|
"""Test get_max_resources_from_cluster_config method.
|
|
|
|
This test verifies that the method correctly:
|
|
1. Always includes CPU/GPU/TPU in the results
|
|
2. Discovers additional resource types from node_group_configs and max_resources
|
|
3. Calculates maximum values for each resource type
|
|
4. Handles edge cases like empty configs, zero counts, unlimited resources
|
|
5. Supports resource types beyond CPU/GPU/TPU
|
|
"""
|
|
ray.init(num_cpus=1)
|
|
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
|
|
|
|
gcs_client.report_cluster_config(cluster_config.SerializeToString())
|
|
max_resources = ray._private.state.state.get_max_resources_from_cluster_config()
|
|
|
|
assert (
|
|
max_resources == expected_resources
|
|
), f"{description}\nExpected: {expected_resources}\nActual: {max_resources}"
|
|
|
|
|
|
def test_get_draining_nodes(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
cluster.add_node()
|
|
ray.init(address=cluster.address)
|
|
cluster.add_node(resources={"worker1": 1})
|
|
cluster.add_node(resources={"worker2": 1})
|
|
cluster.wait_for_nodes()
|
|
|
|
@ray.remote
|
|
def get_node_id():
|
|
return ray.get_runtime_context().get_node_id()
|
|
|
|
worker1_node_id = ray.get(get_node_id.options(resources={"worker1": 1}).remote())
|
|
worker2_node_id = ray.get(get_node_id.options(resources={"worker2": 1}).remote())
|
|
|
|
# Initially there is no draining node.
|
|
assert ray._private.state.state.get_draining_nodes() == {}
|
|
|
|
@ray.remote
|
|
class Actor:
|
|
def ping(self):
|
|
pass
|
|
|
|
actor1 = Actor.options(num_cpus=1, resources={"worker1": 1}).remote()
|
|
actor2 = Actor.options(num_cpus=1, resources={"worker2": 1}).remote()
|
|
ray.get(actor1.ping.remote())
|
|
ray.get(actor2.ping.remote())
|
|
|
|
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
|
|
|
|
# Drain the worker nodes.
|
|
is_accepted, _ = gcs_client.drain_node(
|
|
worker1_node_id,
|
|
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
|
|
"preemption",
|
|
2**63 - 2,
|
|
)
|
|
assert is_accepted
|
|
|
|
is_accepted, _ = gcs_client.drain_node(
|
|
worker2_node_id,
|
|
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
|
|
"preemption",
|
|
0,
|
|
)
|
|
assert is_accepted
|
|
|
|
def get_draining_nodes_check():
|
|
draining_nodes = ray._private.state.state.get_draining_nodes()
|
|
if (
|
|
draining_nodes[worker1_node_id] == (2**63 - 2)
|
|
and draining_nodes[worker2_node_id] == 0
|
|
):
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
wait_for_condition(get_draining_nodes_check)
|
|
|
|
# Kill the actors running on the draining worker nodes so
|
|
# that the worker nodes become idle and can be drained.
|
|
ray.kill(actor1)
|
|
ray.kill(actor2)
|
|
|
|
wait_for_condition(lambda: ray._private.state.state.get_draining_nodes() == {})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|