688 lines
22 KiB
Python
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__]))
|