wehub-resource-sync
341 lines
10 KiB
Python
341 lines
10 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import SignalActor, wait_for_condition
|
|
from ray._private.test_utils import (
|
|
RPC_FAILURE_MAP,
|
|
RPC_FAILURE_TYPES,
|
|
)
|
|
from ray.core.generated import autoscaler_pb2
|
|
from ray.exceptions import GetTimeoutError, TaskCancelledError
|
|
from ray.util.placement_group import placement_group, remove_placement_group
|
|
from ray.util.scheduling_strategies import (
|
|
NodeAffinitySchedulingStrategy,
|
|
PlacementGroupSchedulingStrategy,
|
|
)
|
|
|
|
import psutil
|
|
|
|
|
|
@pytest.mark.parametrize("deterministic_failure", RPC_FAILURE_TYPES)
|
|
def test_request_worker_lease_idempotent(
|
|
monkeypatch, shutdown_only, deterministic_failure, ray_start_cluster
|
|
):
|
|
failure = RPC_FAILURE_MAP[deterministic_failure].copy()
|
|
failure["num_failures"] = 1
|
|
monkeypatch.setenv(
|
|
"RAY_testing_rpc_failure",
|
|
json.dumps({"NodeManagerService.grpc_client.RequestWorkerLease": failure}),
|
|
)
|
|
|
|
@ray.remote
|
|
def simple_task_1():
|
|
return 0
|
|
|
|
@ray.remote
|
|
def simple_task_2():
|
|
return 1
|
|
|
|
# Spin up a two-node cluster where we're targeting scheduling on the
|
|
# remote node via NodeAffinitySchedulingStrategy to test remote RequestWorkerLease
|
|
# calls.
|
|
cluster = ray_start_cluster
|
|
remote_node = cluster.add_node(num_cpus=1)
|
|
|
|
result_ref1 = simple_task_1.options(
|
|
scheduling_strategy=NodeAffinitySchedulingStrategy(
|
|
node_id=remote_node.node_id, soft=False
|
|
)
|
|
).remote()
|
|
result_ref2 = simple_task_2.options(
|
|
scheduling_strategy=NodeAffinitySchedulingStrategy(
|
|
node_id=remote_node.node_id, soft=False
|
|
)
|
|
).remote()
|
|
|
|
assert ray.get([result_ref1, result_ref2]) == [0, 1]
|
|
|
|
|
|
def test_drain_node_idempotent(monkeypatch, shutdown_only, ray_start_cluster):
|
|
# NOTE: not testing response failure since the node is already marked as draining and shuts down gracefully.
|
|
monkeypatch.setenv(
|
|
"RAY_testing_rpc_failure",
|
|
json.dumps(
|
|
{
|
|
"NodeManagerService.grpc_client.DrainRaylet": {
|
|
"num_failures": 1,
|
|
"req_failure_prob": 100,
|
|
"resp_failure_prob": 0,
|
|
"in_flight_failure_prob": 0,
|
|
}
|
|
}
|
|
),
|
|
)
|
|
|
|
cluster = ray_start_cluster
|
|
worker_node = cluster.add_node(num_cpus=1)
|
|
ray.init(address=cluster.address)
|
|
|
|
worker_node_id = worker_node.node_id
|
|
|
|
gcs_client = ray._raylet.GcsClient(address=cluster.address)
|
|
|
|
is_accepted = gcs_client.drain_node(
|
|
worker_node_id,
|
|
autoscaler_pb2.DrainNodeReason.DRAIN_NODE_REASON_IDLE_TERMINATION,
|
|
"Test drain",
|
|
0,
|
|
)
|
|
assert is_accepted
|
|
|
|
# After drain is accepted on an idle node since no tasks are running nor primary objects kept
|
|
# on that raylet, it should be marked idle and gracefully shut down.
|
|
def node_is_dead():
|
|
nodes = ray.nodes()
|
|
for node in nodes:
|
|
if node["NodeID"] == worker_node_id:
|
|
return not node["Alive"]
|
|
return True
|
|
|
|
wait_for_condition(node_is_dead, timeout=1)
|
|
|
|
|
|
# Bundles can be leaked if the gcs dies before the RemovePlacementGroupBundles RPCs are
|
|
# propagated to all the raylets. Since this is inherently racy, we block RemovePlacementGroupBundles RPCs
|
|
# from ever succeeding to make this test deterministic.
|
|
@pytest.fixture
|
|
def inject_release_unused_bundles_rpc_failure(monkeypatch, request):
|
|
deterministic_failure = request.param
|
|
failure = RPC_FAILURE_MAP[deterministic_failure].copy()
|
|
failure["num_failures"] = 1
|
|
monkeypatch.setenv(
|
|
"RAY_testing_rpc_failure",
|
|
json.dumps(
|
|
{
|
|
"NodeManagerService.grpc_client.ReleaseUnusedBundles": failure,
|
|
"NodeManagerService.grpc_client.RemovePlacementGroupBundles": {
|
|
"num_failures": -1,
|
|
"req_failure_prob": 100,
|
|
"resp_failure_prob": 0,
|
|
"in_flight_failure_prob": 0,
|
|
},
|
|
}
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"inject_release_unused_bundles_rpc_failure",
|
|
RPC_FAILURE_TYPES,
|
|
indirect=True,
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"ray_start_cluster_head_with_external_redis",
|
|
[{"num_cpus": 1}],
|
|
indirect=True,
|
|
)
|
|
def test_release_unused_bundles_idempotent(
|
|
inject_release_unused_bundles_rpc_failure,
|
|
ray_start_cluster_head_with_external_redis,
|
|
):
|
|
cluster = ray_start_cluster_head_with_external_redis
|
|
|
|
@ray.remote(num_cpus=1)
|
|
def task():
|
|
return "success"
|
|
|
|
pg = placement_group(name="test_pg", strategy="PACK", bundles=[{"CPU": 1}])
|
|
|
|
result_ref = task.options(
|
|
scheduling_strategy=PlacementGroupSchedulingStrategy(
|
|
placement_group=pg,
|
|
placement_group_bundle_index=0,
|
|
)
|
|
).remote()
|
|
assert ray.get(result_ref) == "success"
|
|
|
|
# Remove the placement group. This will trigger RemovePlacementGroupBundles RPCs which need to be blocked
|
|
# for the placement group bundle to be leaked.
|
|
remove_placement_group(pg)
|
|
|
|
cluster.head_node.kill_gcs_server()
|
|
# ReleaseUnusedBundles only triggers after GCS restart to clean up potentially leaked bundles.
|
|
cluster.head_node.start_gcs_server()
|
|
|
|
# If the leaked bundle wasn't cleaned up, this task will hang due to resource unavailability
|
|
result = ray.get(task.remote())
|
|
assert result == "success"
|
|
|
|
|
|
@pytest.fixture
|
|
def inject_notify_gcs_restart_rpc_failure(monkeypatch, request):
|
|
deterministic_failure = request.param
|
|
failure = RPC_FAILURE_MAP[deterministic_failure].copy()
|
|
failure["num_failures"] = 1
|
|
monkeypatch.setenv(
|
|
"RAY_testing_rpc_failure",
|
|
json.dumps({"NodeManagerService.grpc_client.NotifyGCSRestart": failure}),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"inject_notify_gcs_restart_rpc_failure",
|
|
RPC_FAILURE_TYPES,
|
|
indirect=True,
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"ray_start_cluster_head_with_external_redis",
|
|
[
|
|
{
|
|
"_system_config": {
|
|
# Extending the fallback timeout to focus on death
|
|
# notification received from GCS_ACTOR_CHANNEL pubsub
|
|
"timeout_ms_task_wait_for_death_info": 10000,
|
|
}
|
|
}
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_notify_gcs_restart_idempotent(
|
|
inject_notify_gcs_restart_rpc_failure,
|
|
ray_start_cluster_head_with_external_redis,
|
|
):
|
|
cluster = ray_start_cluster_head_with_external_redis
|
|
|
|
@ray.remote(num_cpus=1, max_restarts=0)
|
|
class DummyActor:
|
|
def get_pid(self):
|
|
return psutil.Process().pid
|
|
|
|
def ping(self):
|
|
return "pong"
|
|
|
|
actor = DummyActor.remote()
|
|
ray.get(actor.ping.remote())
|
|
actor_pid = ray.get(actor.get_pid.remote())
|
|
|
|
cluster.head_node.kill_gcs_server()
|
|
cluster.head_node.start_gcs_server()
|
|
|
|
p = psutil.Process(actor_pid)
|
|
p.kill()
|
|
|
|
# If the actor death notification is not received from the GCS pubsub, this will timeout since
|
|
# the fallback via wait_for_death_info_tasks in the actor task submitter will never trigger
|
|
# since it's set to 10 seconds.
|
|
with pytest.raises(ray.exceptions.RayActorError):
|
|
ray.get(actor.ping.remote(), timeout=5)
|
|
|
|
|
|
def test_kill_local_actor_rpc_retry_and_idempotency(monkeypatch, shutdown_only):
|
|
"""Test that KillLocalActor RPC retries work correctly and guarantee actor death.
|
|
Not testing response since the actor is killed either way.
|
|
"""
|
|
|
|
monkeypatch.setenv(
|
|
"RAY_testing_rpc_failure",
|
|
json.dumps(
|
|
{
|
|
"NodeManagerService.grpc_client.KillLocalActor": {
|
|
"num_failures": 1,
|
|
"req_failure_prob": 100,
|
|
"resp_failure_prob": 0,
|
|
"in_flight_failure_prob": 0,
|
|
}
|
|
}
|
|
),
|
|
)
|
|
|
|
ray.init()
|
|
|
|
@ray.remote
|
|
class SimpleActor:
|
|
def ping(self):
|
|
return "pong"
|
|
|
|
def get_pid(self):
|
|
return os.getpid()
|
|
|
|
actor = SimpleActor.remote()
|
|
|
|
result = ray.get(actor.ping.remote())
|
|
assert result == "pong"
|
|
|
|
worker_pid = ray.get(actor.get_pid.remote())
|
|
|
|
# NOTE: checking the process is still alive rather than checking the actor state from the GCS
|
|
# since as long as KillActor is sent the GCS will mark the actor as dead even though it may not actually be
|
|
assert psutil.pid_exists(worker_pid)
|
|
|
|
ray.kill(actor)
|
|
|
|
def verify_process_killed():
|
|
return not psutil.pid_exists(worker_pid)
|
|
|
|
wait_for_condition(verify_process_killed, timeout=30)
|
|
|
|
|
|
@pytest.fixture
|
|
def inject_cancel_local_task_rpc_failure(monkeypatch, request):
|
|
failure = RPC_FAILURE_MAP[request.param].copy()
|
|
failure["num_failures"] = 1
|
|
monkeypatch.setenv(
|
|
"RAY_testing_rpc_failure",
|
|
json.dumps(
|
|
{
|
|
"NodeManagerService.grpc_client.CancelLocalTask": failure,
|
|
}
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"inject_cancel_local_task_rpc_failure", RPC_FAILURE_TYPES, indirect=True
|
|
)
|
|
@pytest.mark.parametrize("force_kill", [True, False])
|
|
def test_cancel_local_task_rpc_retry_and_idempotency(
|
|
inject_cancel_local_task_rpc_failure, force_kill, shutdown_only
|
|
):
|
|
"""Test that CancelLocalTask RPC retries work correctly.
|
|
|
|
Verify that the RPC is idempotent when network failures occur.
|
|
When force_kill=True, verify the worker process is actually killed using psutil.
|
|
"""
|
|
ray.init(num_cpus=1)
|
|
signaler = SignalActor.remote()
|
|
|
|
@ray.remote(num_cpus=1)
|
|
def get_pid():
|
|
return os.getpid()
|
|
|
|
@ray.remote(num_cpus=1)
|
|
def blocking_task():
|
|
return ray.get(signaler.wait.remote())
|
|
|
|
worker_pid = ray.get(get_pid.remote())
|
|
|
|
blocking_ref = blocking_task.remote()
|
|
|
|
with pytest.raises(GetTimeoutError):
|
|
ray.get(blocking_ref, timeout=1)
|
|
|
|
ray.cancel(blocking_ref, force=force_kill)
|
|
|
|
with pytest.raises(TaskCancelledError):
|
|
ray.get(blocking_ref, timeout=10)
|
|
|
|
if force_kill:
|
|
|
|
def verify_process_killed():
|
|
return not psutil.pid_exists(worker_pid)
|
|
|
|
wait_for_condition(verify_process_killed, timeout=30)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|