chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,591 @@
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import sys
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from typing import Set
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray._common.test_utils import SignalActor, wait_for_condition
|
||||
from ray.exceptions import TaskCancelledError
|
||||
from ray.util.state import list_tasks
|
||||
|
||||
|
||||
def test_input_validation(shutdown_only):
|
||||
# Verify force=True is not working.
|
||||
@ray.remote
|
||||
class A:
|
||||
async def f(self):
|
||||
pass
|
||||
|
||||
a = A.remote()
|
||||
with pytest.raises(ValueError, match="force=True is not supported"):
|
||||
ray.cancel(a.f.remote(), force=True)
|
||||
|
||||
|
||||
def test_async_actor_cancel(shutdown_only):
|
||||
"""
|
||||
Test async actor task is canceled and
|
||||
asyncio.CancelledError is raised within a task.
|
||||
"""
|
||||
ray.init(num_cpus=1)
|
||||
|
||||
@ray.remote
|
||||
class VerifyActor:
|
||||
def __init__(self):
|
||||
self.called = False
|
||||
self.running = False
|
||||
|
||||
def called(self):
|
||||
self.called = True
|
||||
|
||||
def set_running(self):
|
||||
self.running = True
|
||||
|
||||
def is_called(self):
|
||||
return self.called
|
||||
|
||||
def is_running(self):
|
||||
return self.running
|
||||
|
||||
def reset(self):
|
||||
self.called = False
|
||||
self.running = False
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
async def f(self, verify_actor):
|
||||
try:
|
||||
ray.get(verify_actor.set_running.remote())
|
||||
await asyncio.sleep(10)
|
||||
except asyncio.CancelledError:
|
||||
# It is False until this except block is finished.
|
||||
assert not asyncio.current_task().cancelled()
|
||||
ray.get(verify_actor.called.remote())
|
||||
raise
|
||||
except Exception:
|
||||
return True
|
||||
return True
|
||||
|
||||
v = VerifyActor.remote()
|
||||
a = Actor.remote()
|
||||
|
||||
for i in range(50):
|
||||
ref = a.f.remote(v)
|
||||
wait_for_condition(lambda: ray.get(v.is_running.remote()))
|
||||
ray.cancel(ref)
|
||||
|
||||
with pytest.raises(ray.exceptions.TaskCancelledError, match="was cancelled"):
|
||||
ray.get(ref)
|
||||
|
||||
# Verify asyncio.CancelledError is raised from the actor task.
|
||||
assert ray.get(v.is_running.remote())
|
||||
assert ray.get(v.is_called.remote())
|
||||
ray.get(v.reset.remote())
|
||||
|
||||
|
||||
def test_async_actor_client_side_cancel(ray_start_cluster):
|
||||
"""
|
||||
Test a task is cancelled while it is queued on a client side.
|
||||
It should raise ray.exceptions.TaskCancelledError.
|
||||
"""
|
||||
cluster = ray_start_cluster
|
||||
cluster.add_node(num_cpus=0)
|
||||
ray.init(address=cluster.address)
|
||||
|
||||
@ray.remote(num_cpus=1)
|
||||
class Actor:
|
||||
def __init__(self):
|
||||
self.f_called = False
|
||||
|
||||
async def g(self, ref):
|
||||
await asyncio.sleep(30)
|
||||
|
||||
async def f(self):
|
||||
self.f_called = True
|
||||
await asyncio.sleep(5)
|
||||
|
||||
def is_f_called(self):
|
||||
return self.f_called
|
||||
|
||||
@ray.remote
|
||||
def f():
|
||||
time.sleep(100)
|
||||
|
||||
# Test the case where a task is queued on a client side.
|
||||
# Tasks are not sent until actor is created.
|
||||
a = Actor.remote()
|
||||
ref = a.f.remote()
|
||||
ray.cancel(ref)
|
||||
with pytest.raises(TaskCancelledError):
|
||||
ray.get(ref)
|
||||
|
||||
cluster.add_node(num_cpus=1)
|
||||
assert not ray.get(a.is_f_called.remote())
|
||||
|
||||
# Test the case where it is canceled before dependencies
|
||||
# are resolved.
|
||||
a = Actor.remote()
|
||||
ref_dep_not_resolved = a.g.remote(f.remote())
|
||||
ray.cancel(ref_dep_not_resolved)
|
||||
with pytest.raises(TaskCancelledError):
|
||||
ray.get(ref_dep_not_resolved)
|
||||
|
||||
|
||||
def test_async_actor_server_side_cancel(shutdown_only):
|
||||
"""
|
||||
Test Cancelation when a task is queued on a server side.
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
async def f(self):
|
||||
await asyncio.sleep(5)
|
||||
|
||||
async def g(self):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
a = Actor.options(max_concurrency=1).remote()
|
||||
ray.get(a.__ray_ready__.remote())
|
||||
ref = a.f.remote() # noqa
|
||||
# Queued on a server side.
|
||||
# Task should not be executed at all.
|
||||
refs = [a.g.remote() for _ in range(100)]
|
||||
wait_for_condition(
|
||||
lambda: len(
|
||||
list_tasks(
|
||||
filters=[
|
||||
("name", "=", "Actor.g"),
|
||||
("STATE", "=", "PENDING_ACTOR_TASK_ORDERING_OR_CONCURRENCY"),
|
||||
]
|
||||
)
|
||||
)
|
||||
== 100
|
||||
)
|
||||
|
||||
for ref in refs:
|
||||
ray.cancel(ref)
|
||||
tasks = list_tasks(filters=[("name", "=", "Actor.g")])
|
||||
|
||||
for ref in refs:
|
||||
with pytest.raises(TaskCancelledError, match=ref.task_id().hex()):
|
||||
ray.get(ref)
|
||||
|
||||
# Verify the task is submitted to the worker and never executed
|
||||
for task in tasks:
|
||||
assert task.state == "PENDING_ACTOR_TASK_ORDERING_OR_CONCURRENCY"
|
||||
|
||||
|
||||
def test_async_actor_cancel_after_task_finishes(shutdown_only):
|
||||
@ray.remote
|
||||
class Actor:
|
||||
async def f(self):
|
||||
await asyncio.sleep(5)
|
||||
|
||||
async def empty(self):
|
||||
pass
|
||||
|
||||
# Cancel after task finishes
|
||||
a = Actor.options(max_concurrency=1).remote()
|
||||
ref = a.empty.remote()
|
||||
ref2 = a.empty.remote()
|
||||
ray.get([ref, ref2])
|
||||
ray.cancel(ref)
|
||||
ray.cancel(ref2)
|
||||
# Exceptions shouldn't be raised.
|
||||
ray.get([ref, ref2])
|
||||
|
||||
|
||||
def test_async_actor_cancel_restart(ray_start_cluster, monkeypatch):
|
||||
"""
|
||||
Verify a cancelation works if actor is restarted.
|
||||
"""
|
||||
with monkeypatch.context() as m:
|
||||
# This will slow down the cancelation RPC so that
|
||||
# cancel won't succeed until a node is killed.
|
||||
m.setenv(
|
||||
"RAY_testing_asio_delay_us",
|
||||
"CoreWorkerService.grpc_server.CancelTask=3000000:3000000",
|
||||
)
|
||||
cluster = ray_start_cluster
|
||||
cluster.add_node(num_cpus=0)
|
||||
ray.init(address=cluster.address)
|
||||
node = cluster.add_node(num_cpus=1)
|
||||
|
||||
@ray.remote(num_cpus=1, max_restarts=-1, max_task_retries=-1)
|
||||
class Actor:
|
||||
async def f(self):
|
||||
await asyncio.sleep(10)
|
||||
|
||||
a = Actor.remote()
|
||||
ref = a.f.remote()
|
||||
# This guarantees that a.f.remote() is executed
|
||||
ray.get(a.__ray_ready__.remote())
|
||||
ray.cancel(ref)
|
||||
cluster.remove_node(node)
|
||||
r, ur = ray.wait([ref])
|
||||
# When cancel is called, the task won't be retried anymore.
|
||||
# It will raise TaskCancelledError.
|
||||
with pytest.raises(ray.exceptions.TaskCancelledError):
|
||||
ray.get(ref)
|
||||
|
||||
# This will restart actor, but task won't be retried.
|
||||
cluster.add_node(num_cpus=1)
|
||||
# Verify actor is restarted. f should be retried
|
||||
ray.get(a.__ray_ready__.remote())
|
||||
with pytest.raises(ray.exceptions.TaskCancelledError):
|
||||
ray.get(ref)
|
||||
|
||||
|
||||
def test_remote_cancel(ray_start_regular):
|
||||
@ray.remote
|
||||
class Actor:
|
||||
async def sleep(self):
|
||||
await asyncio.sleep(1000)
|
||||
|
||||
@ray.remote
|
||||
def f(refs):
|
||||
ref = refs[0]
|
||||
ray.cancel(ref)
|
||||
|
||||
a = Actor.remote()
|
||||
sleep_ref = a.sleep.remote()
|
||||
wait_for_condition(lambda: list_tasks(filters=[("name", "=", "Actor.sleep")]))
|
||||
ref = f.remote([sleep_ref]) # noqa
|
||||
|
||||
with pytest.raises(ray.exceptions.TaskCancelledError):
|
||||
ray.get(sleep_ref)
|
||||
|
||||
|
||||
def test_cancel_recursive_tree(shutdown_only):
|
||||
"""Verify recursive cancel works for tree-nested tasks.
|
||||
|
||||
Task A -> Task B
|
||||
-> Task C
|
||||
"""
|
||||
ray.init(num_cpus=16)
|
||||
|
||||
# Test the tree structure.
|
||||
@ray.remote
|
||||
def child():
|
||||
for _ in range(5):
|
||||
time.sleep(1)
|
||||
return True
|
||||
|
||||
@ray.remote
|
||||
class ChildActor:
|
||||
async def child(self):
|
||||
await asyncio.sleep(5)
|
||||
return True
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def __init__(self):
|
||||
self.children_refs = defaultdict(list)
|
||||
|
||||
def get_children_refs(self, task_id):
|
||||
return self.children_refs[task_id]
|
||||
|
||||
async def run(self, child_actor, sig):
|
||||
ref1 = child.remote()
|
||||
ref2 = child_actor.child.remote()
|
||||
task_id = ray.get_runtime_context().get_task_id()
|
||||
self.children_refs[task_id].append(ref1)
|
||||
self.children_refs[task_id].append(ref2)
|
||||
|
||||
await sig.wait.remote()
|
||||
await ref1
|
||||
await ref2
|
||||
|
||||
sig = SignalActor.remote()
|
||||
child_actor = ChildActor.remote()
|
||||
a = Actor.remote()
|
||||
ray.get(a.__ray_ready__.remote())
|
||||
|
||||
"""
|
||||
Test the basic case.
|
||||
"""
|
||||
run_ref = a.run.remote(child_actor, sig)
|
||||
task_id = run_ref.task_id().hex()
|
||||
wait_for_condition(
|
||||
lambda: list_tasks(filters=[("task_id", "=", task_id)])[0].state == "RUNNING",
|
||||
timeout=20,
|
||||
)
|
||||
ray.cancel(run_ref, recursive=True)
|
||||
ray.get(sig.send.remote())
|
||||
|
||||
children_refs = ray.get(a.get_children_refs.remote(task_id))
|
||||
|
||||
for ref in children_refs + [run_ref]:
|
||||
with pytest.raises(ray.exceptions.TaskCancelledError):
|
||||
ray.get(ref)
|
||||
|
||||
"""
|
||||
Test recursive = False
|
||||
"""
|
||||
run_ref = a.run.remote(child_actor, sig)
|
||||
task_id = run_ref.task_id().hex()
|
||||
wait_for_condition(
|
||||
lambda: list_tasks(filters=[("task_id", "=", task_id)])[0].state == "RUNNING",
|
||||
timeout=20,
|
||||
)
|
||||
ray.cancel(run_ref, recursive=False)
|
||||
ray.get(sig.send.remote())
|
||||
|
||||
children_refs = ray.get(a.get_children_refs.remote(task_id))
|
||||
|
||||
for ref in children_refs:
|
||||
assert ray.get(ref)
|
||||
|
||||
with pytest.raises(ray.exceptions.TaskCancelledError):
|
||||
ray.get(run_ref)
|
||||
|
||||
"""
|
||||
Test concurrent cases.
|
||||
"""
|
||||
run_refs = [a.run.remote(ChildActor.remote(), sig) for _ in range(10)]
|
||||
task_ids = []
|
||||
|
||||
for i, run_ref in enumerate(run_refs):
|
||||
task_id = run_ref.task_id().hex()
|
||||
task_ids.append(task_id)
|
||||
wait_for_condition(
|
||||
lambda task_id=task_id: list_tasks(filters=[("task_id", "=", task_id)])[
|
||||
0
|
||||
].state
|
||||
== "RUNNING",
|
||||
timeout=20,
|
||||
)
|
||||
children_refs = ray.get(a.get_children_refs.remote(task_id))
|
||||
for child_ref in children_refs:
|
||||
task_id = child_ref.task_id().hex()
|
||||
wait_for_condition(
|
||||
lambda task_id=task_id: list_tasks(filters=[("task_id", "=", task_id)])[
|
||||
0
|
||||
].state
|
||||
== "RUNNING",
|
||||
timeout=20,
|
||||
)
|
||||
recursive = i % 2 == 0
|
||||
ray.cancel(run_ref, recursive=recursive)
|
||||
|
||||
ray.get(sig.send.remote())
|
||||
|
||||
for i, task_id in enumerate(task_ids):
|
||||
children_refs = ray.get(a.get_children_refs.remote(task_id))
|
||||
|
||||
if i % 2 == 0:
|
||||
for ref in children_refs:
|
||||
with pytest.raises(ray.exceptions.TaskCancelledError):
|
||||
ray.get(ref)
|
||||
else:
|
||||
for ref in children_refs:
|
||||
assert ray.get(ref)
|
||||
|
||||
with pytest.raises(ray.exceptions.TaskCancelledError):
|
||||
ray.get(run_refs[i])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("recursive", [True, False])
|
||||
def test_cancel_recursive_chain(shutdown_only, recursive):
|
||||
@ray.remote
|
||||
class RecursiveActor:
|
||||
def __init__(self, child=None):
|
||||
self.child = child
|
||||
self.chlid_ref = None
|
||||
|
||||
async def run(self, sig):
|
||||
if self.child is None:
|
||||
await sig.wait.remote()
|
||||
return True
|
||||
|
||||
ref = self.child.run.remote(sig)
|
||||
self.child_ref = ref
|
||||
return await ref
|
||||
|
||||
def get_child_ref(self):
|
||||
return self.child_ref
|
||||
|
||||
sig = SignalActor.remote()
|
||||
r1 = RecursiveActor.remote()
|
||||
r2 = RecursiveActor.remote(r1)
|
||||
r3 = RecursiveActor.remote(r2)
|
||||
r4 = RecursiveActor.remote(r3)
|
||||
|
||||
ref = r4.run.remote(sig)
|
||||
ray.get(r4.__ray_ready__.remote())
|
||||
wait_for_condition(
|
||||
lambda: len(list_tasks(filters=[("name", "=", "RecursiveActor.run")])) == 4
|
||||
)
|
||||
ray.cancel(ref, recursive=recursive)
|
||||
ray.get(sig.send.remote())
|
||||
|
||||
if recursive:
|
||||
with pytest.raises(ray.exceptions.TaskCancelledError):
|
||||
ray.get(ref)
|
||||
with pytest.raises(ray.exceptions.TaskCancelledError):
|
||||
ray.get(ray.get(r4.get_child_ref.remote()))
|
||||
with pytest.raises(ray.exceptions.TaskCancelledError):
|
||||
ray.get(ray.get(r3.get_child_ref.remote()))
|
||||
with pytest.raises(ray.exceptions.TaskCancelledError):
|
||||
ray.get(ray.get(r2.get_child_ref.remote()))
|
||||
else:
|
||||
assert ray.get(ray.get(r2.get_child_ref.remote()))
|
||||
assert ray.get(ray.get(r3.get_child_ref.remote()))
|
||||
assert ray.get(ray.get(r4.get_child_ref.remote()))
|
||||
|
||||
with pytest.raises(ray.exceptions.TaskCancelledError):
|
||||
ray.get(ref)
|
||||
|
||||
|
||||
def test_concurrent_submission_and_cancellation(shutdown_only):
|
||||
"""Test submitting and then cancelling many tasks concurrently.
|
||||
|
||||
This is a regression test for race conditions such as:
|
||||
https://github.com/ray-project/ray/issues/52628.
|
||||
"""
|
||||
NUM_TASKS = 2500
|
||||
|
||||
@ray.remote(num_cpus=0)
|
||||
class Worker:
|
||||
async def sleep(self, i: int):
|
||||
# NOTE: all tasks should be cancelled, so this won't actually sleep for the
|
||||
# full duration if the test is passing.
|
||||
await asyncio.sleep(30)
|
||||
|
||||
worker = Worker.remote()
|
||||
|
||||
# Submit many tasks in parallel to cause queueing on the caller and receiver.
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=NUM_TASKS) as executor:
|
||||
futures = [executor.submit(worker.sleep.remote, i) for i in range(NUM_TASKS)]
|
||||
refs = [f.result() for f in concurrent.futures.as_completed(futures)]
|
||||
|
||||
# Cancel the tasks in reverse order of submission.
|
||||
for ref in reversed(refs):
|
||||
ray.cancel(ref)
|
||||
|
||||
# Check that all tasks were successfully cancelled (none ran to completion).
|
||||
for ref in refs:
|
||||
with pytest.raises(ray.exceptions.TaskCancelledError):
|
||||
ray.get(ref)
|
||||
|
||||
|
||||
def test_is_canceled_sync_actor_task(shutdown_only):
|
||||
"""Test that is_canceled() works correctly for sync actor tasks."""
|
||||
|
||||
signal_actor = SignalActor.remote()
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def __init__(self):
|
||||
self._was_canceled = False
|
||||
|
||||
def wait_until_canceled(self):
|
||||
ray.get(signal_actor.wait.remote())
|
||||
wait_for_condition(lambda: ray.get_runtime_context().is_canceled())
|
||||
self._was_canceled = True
|
||||
|
||||
def was_canceled(self) -> bool:
|
||||
return self._was_canceled
|
||||
|
||||
a = Actor.remote()
|
||||
ref = a.wait_until_canceled.remote()
|
||||
|
||||
# Wait for the task to be actively waiting on the signal.
|
||||
wait_for_condition(lambda: ray.get(signal_actor.cur_num_waiters.remote()) == 1)
|
||||
|
||||
# Cancel the task while it's blocked on the signal.
|
||||
ray.cancel(ref, recursive=False)
|
||||
|
||||
# Now signal the task to unblock. The task result should be `TaskCancelledError`.
|
||||
ray.get(signal_actor.send.remote())
|
||||
with pytest.raises(TaskCancelledError):
|
||||
ray.get(ref)
|
||||
|
||||
# Check that `is_canceled` was set correctly.
|
||||
assert ray.get(a.was_canceled.remote())
|
||||
|
||||
|
||||
def test_is_canceled_concurrent_actor_task(shutdown_only):
|
||||
"""Test that is_canceled() works correctly for concurrent actor tasks."""
|
||||
|
||||
signal_actor = SignalActor.remote()
|
||||
|
||||
@ray.remote
|
||||
class ConcurrentActor:
|
||||
def __init__(self):
|
||||
self._canceled_task_indices = set()
|
||||
|
||||
def task_with_cancel_check(self, task_index: int, expect_canceled: bool):
|
||||
ray.get(signal_actor.wait.remote())
|
||||
|
||||
if expect_canceled:
|
||||
wait_for_condition(lambda: ray.get_runtime_context().is_canceled())
|
||||
self._canceled_task_indices.add(task_index)
|
||||
|
||||
return task_index
|
||||
|
||||
def get_canceled_task_indices(self) -> Set[int]:
|
||||
return self._canceled_task_indices
|
||||
|
||||
actor = ConcurrentActor.options(max_concurrency=3).remote()
|
||||
|
||||
# Submit multiple tasks concurrently. Only task_index=1 will be canceled.
|
||||
refs = [actor.task_with_cancel_check.remote(i, i == 1) for i in range(3)]
|
||||
|
||||
# Wait for all tasks to be running (waiting on the signal).
|
||||
wait_for_condition(lambda: ray.get(signal_actor.cur_num_waiters.remote()) == 3)
|
||||
|
||||
# Cancel task_index=1.
|
||||
ray.cancel(refs[1], recursive=False)
|
||||
|
||||
# Send signal to unblock all tasks.
|
||||
ray.get(signal_actor.send.remote())
|
||||
|
||||
# The canceled task should raise TaskCancelledError.
|
||||
with pytest.raises(TaskCancelledError):
|
||||
ray.get(refs[1])
|
||||
|
||||
# The other tasks should complete normally.
|
||||
assert ray.get([refs[0], refs[2]]) == [0, 2]
|
||||
|
||||
# Verify that `is_canceled` was propagated for task_index=1.
|
||||
assert ray.get(actor.get_canceled_task_indices.remote()) == {1}
|
||||
|
||||
|
||||
def test_is_canceled_not_supported_in_async_actor(shutdown_only):
|
||||
"""Test is_canceled() for async actors."""
|
||||
|
||||
@ray.remote
|
||||
class AsyncActor:
|
||||
def __init__(self):
|
||||
self.is_canceled = False
|
||||
|
||||
async def async_task(self):
|
||||
# is_canceled() doesn't work for async actors
|
||||
if ray.get_runtime_context().is_canceled():
|
||||
self.is_canceled = True
|
||||
return "canceled"
|
||||
return "completed"
|
||||
|
||||
def is_canceled(self):
|
||||
return self.is_canceled
|
||||
|
||||
actor = AsyncActor.remote()
|
||||
ref = actor.async_task.remote()
|
||||
|
||||
# is_canceled() is not supported for async actors
|
||||
with pytest.raises(
|
||||
RuntimeError, match="This method is not supported in an async actor."
|
||||
):
|
||||
ray.get(ref)
|
||||
|
||||
# Verify the state for async actor does NOT change as there's no graceful
|
||||
# termination for async actor task
|
||||
assert not ray.get(actor.is_canceled.remote())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
Reference in New Issue
Block a user