Files
2026-07-13 13:17:40 +08:00

688 lines
22 KiB
Python

import sys
import time
from collections import Counter
import pytest
import ray
from ray._common.test_utils import SignalActor, wait_for_condition
from ray._raylet import GcsClient
from ray.core.generated import autoscaler_pb2, common_pb2
from ray.util.scheduling_strategies import (
NodeAffinitySchedulingStrategy,
PlacementGroupSchedulingStrategy,
)
from ray.util.state import list_tasks
def test_idle_termination(ray_start_cluster):
cluster = ray_start_cluster
head_node = cluster.add_node(resources={"head": 1})
ray.init(address=cluster.address)
worker_node = cluster.add_node(resources={"worker": 1})
cluster.wait_for_nodes()
head_node_id = head_node.node_id
worker_node_id = worker_node.node_id
wait_for_condition(
lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
== {head_node_id, worker_node_id}
)
@ray.remote(num_cpus=1, resources={"worker": 1})
class Actor:
def ping(self):
pass
actor = Actor.remote()
ray.get(actor.ping.remote())
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
# The worker node is not idle so the drain request should be rejected.
is_accepted, rejection_reason_message = gcs_client.drain_node(
worker_node_id,
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_IDLE_TERMINATION"),
"idle for long enough",
2**63 - 1,
)
assert not is_accepted
assert (
"The node to be idle terminated is no longer idle." in rejection_reason_message
)
ray.kill(actor)
def drain_until_accept():
# The worker node is idle now so the drain request should be accepted.
is_accepted, _ = gcs_client.drain_node(
worker_node_id,
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_IDLE_TERMINATION"),
"idle for long enough",
2**63 - 1,
)
return is_accepted
wait_for_condition(drain_until_accept)
wait_for_condition(
lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
== {head_node_id}
)
worker_node = [node for node in ray.nodes() if node["NodeID"] == worker_node_id][0]
assert worker_node["DeathReason"] == common_pb2.NodeDeathInfo.Reason.Value(
"AUTOSCALER_DRAIN_IDLE"
)
assert worker_node["DeathReasonMessage"] == "idle for long enough"
# Draining a dead node is always accepted.
is_accepted, _ = gcs_client.drain_node(
worker_node_id,
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_IDLE_TERMINATION"),
"idle for long enough",
2**63 - 1,
)
assert is_accepted
def test_idle_drain_rejected_while_holding_pinned_object(ray_start_cluster):
# Tests that a node is not drained if it's holding a pinned object
#
# The resource-sync period is widened so the object-store-memory idle signal is NOT
# refreshed between pinning the object and the drain request. The raylet must refresh
# it at the drain decision regardless; without that, the cached signal is stale and
# the drain is wrongly accepted once CPU/lease go idle.
cluster = ray_start_cluster
cluster.add_node(
num_cpus=0,
resources={"head": 1},
_system_config={"raylet_report_resources_period_milliseconds": 120000},
)
ray.init(address=cluster.address)
worker_node = cluster.add_node(num_cpus=1, resources={"worker": 1})
cluster.wait_for_nodes()
worker_node_id = worker_node.node_id
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
# >100 KiB so the return is stored in the worker's plasma (in_plasma), not inlined
# into the owner's reply.
@ray.remote(num_cpus=1, max_retries=0)
def make_object():
return b"x" * (10 * 1024 * 1024)
# Runs on the worker (head has 0 CPU). Keep the ref but do NOT fetch it, so the only
# copy stays pinned in the worker's plasma.
ref = make_object.remote()
ready, _ = ray.wait([ref], num_returns=1, timeout=60, fetch_local=False)
assert ready
def drain_idle():
is_accepted, _ = gcs_client.drain_node(
worker_node_id,
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_IDLE_TERMINATION"),
"idle drain while node holds a referenced object",
2**63 - 1,
)
return is_accepted
# Make idle drain requests during the window where worker lease is returned (CPU +
# worker footprint go idle). The object held by the node should prevent it from draining.
for _ in range(10):
assert not drain_idle(), "idle drain accepted while the node held an object"
time.sleep(0.5)
# The worker survived and the object is still retrievable.
assert worker_node_id in {node["NodeID"] for node in ray.nodes() if node["Alive"]}
assert len(ray.get(ref)) == 10 * 1024 * 1024
def test_preemption(ray_start_cluster):
cluster = ray_start_cluster
head_node = cluster.add_node(resources={"head": 1})
ray.init(address=cluster.address)
worker_node = cluster.add_node(resources={"worker": 1})
cluster.wait_for_nodes()
head_node_id = head_node.node_id
worker_node_id = worker_node.node_id
@ray.remote(num_cpus=1, resources={"worker": 1})
class Actor:
def ping(self):
pass
actor = Actor.remote()
ray.get(actor.ping.remote())
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
with pytest.raises(ray.exceptions.RaySystemError):
# Test invalid draining deadline
gcs_client.drain_node(
worker_node_id,
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
"preemption",
-1,
)
# The worker node is not idle but the drain request should be still accepted.
is_accepted, _ = gcs_client.drain_node(
worker_node_id,
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
"preemption",
2**63 - 1,
)
assert is_accepted
time.sleep(1)
# Worker node should still be alive since it's not idle and cannot be drained.
wait_for_condition(
lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
== {head_node_id, worker_node_id}
)
ray.kill(actor)
wait_for_condition(
lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
== {head_node_id}
)
worker_node = [node for node in ray.nodes() if node["NodeID"] == worker_node_id][0]
assert worker_node["DeathReason"] == common_pb2.NodeDeathInfo.Reason.Value(
"AUTOSCALER_DRAIN_PREEMPTED"
)
assert worker_node["DeathReasonMessage"] == "preemption"
@pytest.mark.parametrize(
"graceful",
[True, False],
)
def test_preemption_after_draining_deadline(monkeypatch, ray_start_cluster, graceful):
monkeypatch.setenv("RAY_health_check_failure_threshold", "3")
monkeypatch.setenv("RAY_health_check_timeout_ms", "100")
monkeypatch.setenv("RAY_health_check_period_ms", "1000")
monkeypatch.setenv("RAY_health_check_initial_delay_ms", "0")
cluster = ray_start_cluster
head_node = cluster.add_node(resources={"head": 1})
ray.init(address=cluster.address)
worker_node = cluster.add_node(resources={"worker": 1})
cluster.wait_for_nodes()
head_node_id = head_node.node_id
worker_node_id = worker_node.node_id
wait_for_condition(
lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
== {head_node_id, worker_node_id}
)
@ray.remote(num_cpus=1, resources={"worker": 1})
class Actor:
def ping(self):
pass
actor = Actor.remote()
ray.get(actor.ping.remote())
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
# The worker node is not idle but the drain request should be still accepted.
is_accepted, _ = gcs_client.drain_node(
worker_node_id,
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
"preemption",
1,
)
assert is_accepted
# Simulate autoscaler terminates the worker node after the draining deadline.
cluster.remove_node(worker_node, graceful)
wait_for_condition(
lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
== {head_node_id},
)
worker_node = [node for node in ray.nodes() if node["NodeID"] == worker_node_id][0]
assert worker_node["DeathReason"] == common_pb2.NodeDeathInfo.Reason.Value(
"AUTOSCALER_DRAIN_PREEMPTED"
)
assert worker_node["DeathReasonMessage"] == "preemption"
def test_node_death_before_draining_deadline(monkeypatch, ray_start_cluster):
monkeypatch.setenv("RAY_health_check_failure_threshold", "3")
monkeypatch.setenv("RAY_health_check_timeout_ms", "100")
monkeypatch.setenv("RAY_health_check_period_ms", "1000")
monkeypatch.setenv("RAY_health_check_initial_delay_ms", "0")
cluster = ray_start_cluster
head_node = cluster.add_node(resources={"head": 1})
ray.init(address=cluster.address)
worker_node = cluster.add_node(resources={"worker": 1})
cluster.wait_for_nodes()
head_node_id = head_node.node_id
worker_node_id = worker_node.node_id
wait_for_condition(
lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
== {head_node_id, worker_node_id}
)
@ray.remote(num_cpus=1, resources={"worker": 1})
class Actor:
def ping(self):
pass
actor = Actor.remote()
ray.get(actor.ping.remote())
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
# The worker node is not idle but the drain request should be still accepted.
is_accepted, _ = gcs_client.drain_node(
worker_node_id,
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
"preemption",
2**63 - 1,
)
assert is_accepted
# Simulate the worker node crashes before the draining deadline.
cluster.remove_node(worker_node, False)
wait_for_condition(
lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
== {head_node_id},
)
# Since worker node failure is detected to be before the draining deadline,
# this is considered as an unexpected termination.
worker_node = [node for node in ray.nodes() if node["NodeID"] == worker_node_id][0]
assert worker_node["DeathReason"] == common_pb2.NodeDeathInfo.Reason.Value(
"UNEXPECTED_TERMINATION"
)
assert (
worker_node["DeathReasonMessage"]
== "health check failed due to missing too many heartbeats"
)
def test_scheduling_placement_groups_during_draining(ray_start_cluster):
"""Test that the draining node is unschedulable for new pgs."""
cluster = ray_start_cluster
node1 = cluster.add_node(num_cpus=1, resources={"node1": 1})
ray.init(address=cluster.address)
node2 = cluster.add_node(num_cpus=1, resources={"node2": 1})
cluster.add_node(num_cpus=2, resources={"node3": 1})
cluster.wait_for_nodes()
node1_id = node1.node_id
node2_id = node2.node_id
node3_id = node2.node_id
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
# The node is idle so the draining request should be accepted.
is_accepted, _ = gcs_client.drain_node(
node3_id,
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
"preemption",
2**63 - 1,
)
assert is_accepted
@ray.remote
def get_node_id():
return ray.get_runtime_context().get_node_id()
# Even though node3 is the best for pack but it's draining
# so the pg should be on node1 and node2
pg = ray.util.placement_group(bundles=[{"CPU": 1}, {"CPU": 1}], strategy="PACK")
{
ray.get(
get_node_id.options(
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=pg,
placement_group_bundle_index=0,
)
).remote()
),
ray.get(
get_node_id.options(
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=pg,
placement_group_bundle_index=1,
)
).remote()
),
} == {node1_id, node2_id}
def test_scheduling_tasks_and_actors_during_draining(ray_start_cluster):
"""Test that the draining node is unschedulable for new tasks and actors."""
cluster = ray_start_cluster
head_node = cluster.add_node(num_cpus=1, resources={"head": 1})
ray.init(address=cluster.address)
worker_node = cluster.add_node(num_cpus=1, resources={"worker": 1})
cluster.wait_for_nodes()
head_node_id = head_node.node_id
worker_node_id = worker_node.node_id
@ray.remote
class Actor:
def ping(self):
pass
actor = Actor.options(num_cpus=0, resources={"worker": 1}).remote()
ray.get(actor.ping.remote())
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
# The worker node is not idle but the drain request should be still accepted.
is_accepted, _ = gcs_client.drain_node(
worker_node_id,
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
"preemption",
2**63 - 1,
)
assert is_accepted
@ray.remote
def get_node_id():
return ray.get_runtime_context().get_node_id()
assert (
ray.get(get_node_id.options(scheduling_strategy="SPREAD").remote())
== head_node_id
)
assert (
ray.get(get_node_id.options(scheduling_strategy="SPREAD").remote())
== head_node_id
)
assert (
ray.get(
get_node_id.options(
scheduling_strategy=NodeAffinitySchedulingStrategy(
worker_node_id, soft=True
)
).remote()
)
== head_node_id
)
with pytest.raises(ray.exceptions.TaskUnschedulableError):
ray.get(
get_node_id.options(
label_selector={ray._raylet.RAY_NODE_ID_KEY: worker_node_id}
).remote()
)
head_actor = Actor.options(num_cpus=1, resources={"head": 1}).remote()
ray.get(head_actor.ping.remote())
obj = get_node_id.remote()
# Cannot run on the draining worker node even though it has resources.
with pytest.raises(ray.exceptions.GetTimeoutError):
ray.get(obj, timeout=2)
ray.kill(head_actor)
ray.get(obj, timeout=2) == head_node_id
@pytest.mark.parametrize(
"graceful",
[False, True],
)
def test_draining_reason(ray_start_cluster, graceful):
cluster = ray_start_cluster
cluster.add_node(num_cpus=1, resources={"node1": 1})
ray.init(
address=cluster.address,
)
node2 = cluster.add_node(num_cpus=1, resources={"node2": 1})
@ray.remote
class Actor:
def ping(self):
pass
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
node2_id = node2.node_id
# Schedule actor
actor = Actor.options(num_cpus=0, resources={"node2": 1}).remote()
ray.get(actor.ping.remote())
drain_reason_message = "testing node preemption."
# Preemption is always accepted.
is_accepted, _ = gcs_client.drain_node(
node2_id,
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
drain_reason_message,
1,
)
assert is_accepted
# Simulate autoscaler terminates the worker node after the draining deadline.
cluster.remove_node(node2, graceful)
def check_actor_died_error():
try:
ray.get(actor.ping.remote())
return False
except ray.exceptions.ActorDiedError as e:
assert e.preempted
if graceful:
assert "The actor died because its node has died." in str(e)
assert "the actor's node was preempted: " + drain_reason_message in str(
e
)
return True
wait_for_condition(check_actor_died_error)
def test_drain_node_actor_restart(ray_start_cluster):
cluster = ray_start_cluster
cluster.add_node(num_cpus=1, resources={"head": 1})
ray.init(address=cluster.address)
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
@ray.remote(max_restarts=1)
class Actor:
def get_node_id(self):
return ray.get_runtime_context().get_node_id()
# Prepare the first worker node for the actor.
cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})
cluster.wait_for_nodes()
actor = Actor.options(num_cpus=0, resources={"worker": 1}).remote()
def actor_started():
node_id = ray.get(actor.get_node_id.remote())
return node_id == cur_worker.node_id
wait_for_condition(actor_started, timeout=5)
# Kill the current worker node.
cluster.remove_node(cur_worker, True)
# Prepare a new worker node for the actor to be restarted on later.
cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})
cluster.wait_for_nodes()
# Make sure the actor is restarted on the new worker node.
# This should be counted into the max_restarts of the actor.
wait_for_condition(actor_started, timeout=5)
# Preemption the current worker node.
is_accepted, _ = gcs_client.drain_node(
cur_worker.node_id,
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
"preemption",
1,
)
assert is_accepted
cluster.remove_node(cur_worker, True)
# Prepare a new worker node for the actor to be restarted on later.
cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})
cluster.wait_for_nodes()
# Make sure the actor is restarted on the new worker node.
# This should not be counted into the max_restarts of the actor because the actor was preempted.
wait_for_condition(actor_started, timeout=5)
# Kill the current worker node.
cluster.remove_node(cur_worker, True)
# Prepare a new worker node, however, the actor should not be restarted on this node, since
# the max_restarts is reached.
cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})
cluster.wait_for_nodes()
# The actor should not be restarted, thus an exception should be raised.
with pytest.raises(RuntimeError):
wait_for_condition(actor_started, timeout=5)
def test_drain_node_task_retry(ray_start_cluster):
cluster = ray_start_cluster
cluster.add_node(num_cpus=1, resources={"head": 100})
ray.init(address=cluster.address)
cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})
cluster.wait_for_nodes()
node_ids = Counter()
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
@ray.remote(resources={"head": 1})
class NodeTracker:
def __init__(self):
self._node_ids = Counter()
def add_node(self, node_id):
self._node_ids.update([node_id])
def nodes(self):
return self._node_ids
@ray.remote(max_retries=1, resources={"worker": 1})
def func(signal, nodes):
node_id = ray.get_runtime_context().get_node_id()
ray.get(nodes.add_node.remote(node_id))
ray.get(signal.wait.remote())
return node_id
signal = SignalActor.options(resources={"head": 1}).remote()
node_tracker = NodeTracker.remote()
r1 = func.remote(signal, node_tracker)
# Verify the first node is added to the counter by the func.remote task.
node_ids.update([cur_worker.node_id])
wait_for_condition(lambda: ray.get(node_tracker.nodes.remote()) == node_ids)
# Remove the current worker node and add a new one to trigger a retry.
cluster.remove_node(cur_worker, True)
cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})
# Verify the second node is added to the counter by the task after a retry.
node_ids.update([cur_worker.node_id])
wait_for_condition(lambda: ray.get(node_tracker.nodes.remote()) == node_ids)
# Preempt the second node and add a new one to trigger a retry.
is_accepted, _ = gcs_client.drain_node(
cur_worker.node_id,
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
"preemption",
1,
)
assert is_accepted
cluster.remove_node(cur_worker, True)
cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})
# Verify the third node is added to the counter after a preemption retry.
node_ids.update([cur_worker.node_id])
wait_for_condition(lambda: ray.get(node_tracker.nodes.remote()) == node_ids)
# Remove the third node and add a new one, but the task should not retry.
cluster.remove_node(cur_worker, True)
cur_worker = cluster.add_node(num_cpus=1, resources={"worker": 1})
# max_retries is reached, the task should fail.
with pytest.raises(ray.exceptions.NodeDiedError):
ray.get(r1)
def test_leases_rescheduling_during_draining(ray_start_cluster):
"""Test that when a node is being drained, leases inside local lease manager
will be cancelled and re-added to the cluster lease manager for rescheduling
instead of being marked as permanently infeasible.
This is regression test for https://github.com/ray-project/ray/pull/57834/
"""
cluster = ray_start_cluster
cluster.add_node(num_cpus=0)
ray.init(address=cluster.address)
worker1 = cluster.add_node(num_cpus=1)
cluster.wait_for_nodes()
gcs_client = GcsClient(address=ray.get_runtime_context().gcs_address)
@ray.remote(num_cpus=1)
class Actor:
def ping(self):
pass
actor = Actor.remote()
ray.get(actor.ping.remote())
@ray.remote(num_cpus=1)
def get_node_id():
return ray.get_runtime_context().get_node_id()
obj_ref = get_node_id.options(name="f1").remote()
def verify_f1_pending_node_assignment():
tasks = list_tasks(filters=[("name", "=", "f1")])
assert len(tasks) == 1
assert tasks[0]["state"] == "PENDING_NODE_ASSIGNMENT"
return True
# f1 should be in the local lease manager of worker1,
# waiting for resource to be available.
wait_for_condition(verify_f1_pending_node_assignment)
is_accepted, _ = gcs_client.drain_node(
worker1.node_id,
autoscaler_pb2.DrainNodeReason.Value("DRAIN_NODE_REASON_PREEMPTION"),
"preemption",
2**63 - 1,
)
assert is_accepted
# The task should be rescheduled on another node.
worker2 = cluster.add_node(num_cpus=1)
assert ray.get(obj_ref) == worker2.node_id
if __name__ == "__main__":
sys.exit(pytest.main(["-sv", __file__]))