wehub-resource-sync
442 lines
12 KiB
Python
442 lines
12 KiB
Python
# coding: utf-8
|
|
import asyncio
|
|
import os
|
|
import sys
|
|
import threading
|
|
import time
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import SignalActor, wait_for_condition
|
|
from ray._private.client_mode_hook import client_mode_should_convert
|
|
from ray._private.test_utils import (
|
|
kill_actor_and_wait_for_failure,
|
|
wait_for_pid_to_exit,
|
|
)
|
|
from ray.util.state import get_actor
|
|
|
|
|
|
def test_asyncio_actor(ray_start_regular_shared):
|
|
@ray.remote
|
|
class AsyncBatcher:
|
|
def __init__(self):
|
|
self.batch = []
|
|
self.event = asyncio.Event()
|
|
|
|
async def add(self, x):
|
|
self.batch.append(x)
|
|
if len(self.batch) >= 3:
|
|
self.event.set()
|
|
else:
|
|
await self.event.wait()
|
|
return sorted(self.batch)
|
|
|
|
a = AsyncBatcher.remote()
|
|
x1 = a.add.remote(1)
|
|
x2 = a.add.remote(2)
|
|
x3 = a.add.remote(3)
|
|
r1 = ray.get(x1)
|
|
r2 = ray.get(x2)
|
|
r3 = ray.get(x3)
|
|
assert r1 == [1, 2, 3]
|
|
assert r1 == r2 == r3
|
|
|
|
|
|
def test_asyncio_actor_same_thread(ray_start_regular_shared):
|
|
@ray.remote
|
|
class Actor:
|
|
def sync_thread_id(self):
|
|
return threading.current_thread().ident
|
|
|
|
async def async_thread_id(self):
|
|
return threading.current_thread().ident
|
|
|
|
a = Actor.remote()
|
|
sync_id, async_id = ray.get([a.sync_thread_id.remote(), a.async_thread_id.remote()])
|
|
assert sync_id == async_id
|
|
|
|
|
|
def test_asyncio_actor_concurrency(ray_start_regular_shared):
|
|
@ray.remote
|
|
class RecordOrder:
|
|
def __init__(self):
|
|
self.history = []
|
|
|
|
async def do_work(self):
|
|
self.history.append("STARTED")
|
|
# Force a context switch
|
|
await asyncio.sleep(0)
|
|
self.history.append("ENDED")
|
|
|
|
def get_history(self):
|
|
return self.history
|
|
|
|
num_calls = 10
|
|
|
|
a = RecordOrder.options(max_concurrency=1).remote()
|
|
ray.get([a.do_work.remote() for _ in range(num_calls)])
|
|
history = ray.get(a.get_history.remote())
|
|
|
|
# We only care about ordered start-end-start-end sequence because
|
|
# coroutines may be executed out of enqueued order.
|
|
answer = []
|
|
for _ in range(num_calls):
|
|
for status in ["STARTED", "ENDED"]:
|
|
answer.append(status)
|
|
|
|
assert history == answer
|
|
|
|
|
|
def test_asyncio_actor_high_concurrency(ray_start_regular_shared):
|
|
# This tests actor can handle concurrency above recursionlimit.
|
|
|
|
@ray.remote
|
|
class AsyncConcurrencyBatcher:
|
|
def __init__(self, batch_size):
|
|
self.batch = []
|
|
self.event = asyncio.Event()
|
|
self.batch_size = batch_size
|
|
|
|
async def add(self, x):
|
|
self.batch.append(x)
|
|
if len(self.batch) >= self.batch_size:
|
|
self.event.set()
|
|
else:
|
|
await self.event.wait()
|
|
return sorted(self.batch)
|
|
|
|
batch_size = sys.getrecursionlimit()
|
|
actor = AsyncConcurrencyBatcher.options(max_concurrency=batch_size * 2).remote(
|
|
batch_size
|
|
)
|
|
result = ray.get([actor.add.remote(i) for i in range(batch_size)])
|
|
assert result[0] == list(range(batch_size))
|
|
assert result[-1] == list(range(batch_size))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asyncio_get(ray_start_regular_shared, event_loop):
|
|
loop = event_loop
|
|
asyncio.set_event_loop(loop)
|
|
loop.set_debug(True)
|
|
|
|
# Test Async Plasma
|
|
@ray.remote
|
|
def task():
|
|
return 1
|
|
|
|
assert await task.remote() == 1
|
|
|
|
@ray.remote
|
|
def task_throws():
|
|
_ = 1 / 0
|
|
|
|
with pytest.raises(ray.exceptions.RayTaskError):
|
|
await task_throws.remote()
|
|
|
|
# Test actor calls.
|
|
str_len = 200 * 1024
|
|
|
|
@ray.remote
|
|
class Actor:
|
|
def echo(self, i):
|
|
return i
|
|
|
|
def big_object(self):
|
|
# 100Kb is the limit for direct call
|
|
return "a" * (str_len)
|
|
|
|
def throw_error(self):
|
|
_ = 1 / 0
|
|
|
|
actor = Actor.remote()
|
|
|
|
assert await actor.echo.remote(2) == 2
|
|
|
|
assert await actor.big_object.remote() == "a" * str_len
|
|
|
|
with pytest.raises(ray.exceptions.RayTaskError):
|
|
await actor.throw_error.remote()
|
|
|
|
# Wrap in Remote Function to work with Ray client.
|
|
kill_actor_ref = ray.remote(kill_actor_and_wait_for_failure).remote(actor)
|
|
ray.get(kill_actor_ref)
|
|
|
|
with pytest.raises(ray.exceptions.RayActorError):
|
|
await actor.echo.remote(1)
|
|
|
|
|
|
def test_asyncio_actor_async_get(ray_start_regular_shared):
|
|
@ray.remote
|
|
def remote_task():
|
|
return 1
|
|
|
|
@ray.remote
|
|
class AsyncGetter:
|
|
async def get(self):
|
|
return await remote_task.remote()
|
|
|
|
async def plasma_get(self, plasma_object):
|
|
return await plasma_object[0]
|
|
|
|
plasma_object = ray.put(2)
|
|
getter = AsyncGetter.remote()
|
|
assert ray.get(getter.get.remote()) == 1
|
|
assert ray.get(getter.plasma_get.remote([plasma_object])) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asyncio_double_await(ray_start_regular_shared):
|
|
# This is a regression test for
|
|
# https://github.com/ray-project/ray/issues/8841
|
|
|
|
signal = SignalActor.remote()
|
|
waiting = signal.wait.remote()
|
|
|
|
future = asyncio.ensure_future(waiting)
|
|
with pytest.raises(asyncio.TimeoutError):
|
|
await asyncio.wait_for(future, timeout=0.1)
|
|
assert future.cancelled()
|
|
|
|
# We are explicitly waiting multiple times here to test asyncio state
|
|
# override.
|
|
await signal.send.remote()
|
|
await waiting
|
|
await waiting
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"ray_start_regular_shared",
|
|
[{"include_dashboard": True}],
|
|
indirect=True,
|
|
)
|
|
async def test_asyncio_exit_actor(ray_start_regular_shared):
|
|
# https://github.com/ray-project/ray/issues/12649
|
|
# The test should just hang without the fix.
|
|
|
|
@ray.remote
|
|
class Actor:
|
|
async def exit(self):
|
|
ray.actor.exit_actor()
|
|
|
|
async def ping(self):
|
|
return os.getpid()
|
|
|
|
async def loop_forever(self):
|
|
while True:
|
|
await asyncio.sleep(5)
|
|
|
|
a = Actor.options(max_task_retries=0).remote()
|
|
pid = ray.get(a.ping.remote())
|
|
a.loop_forever.remote()
|
|
# Make sure exit_actor exits immediately, not once all tasks completed.
|
|
with pytest.raises(ray.exceptions.RayError):
|
|
ray.get(a.exit.remote())
|
|
|
|
# New calls should just error.
|
|
with pytest.raises(ray.exceptions.RayError):
|
|
ray.get(a.ping.remote())
|
|
|
|
# The actor should be dead in the actor table.
|
|
# Using ray task so it works in Ray Client as well.
|
|
@ray.remote
|
|
def check_actor_gone_now():
|
|
def cond():
|
|
return get_actor(id=a._ray_actor_id.hex()).state != "ALIVE"
|
|
|
|
wait_for_condition(cond)
|
|
|
|
ray.get(check_actor_gone_now.remote())
|
|
|
|
# Make sure there is no process leak
|
|
wait_for_pid_to_exit(pid)
|
|
|
|
|
|
def test_asyncio_exit_actor_with_concurrency_group(ray_start_regular_shared):
|
|
@ray.remote(concurrency_groups={"async": 2})
|
|
class Actor:
|
|
async def getpid(self):
|
|
return os.getpid()
|
|
|
|
async def exit(self):
|
|
ray.actor.exit_actor()
|
|
|
|
@ray.method(concurrency_group="async")
|
|
async def loop_forever(self):
|
|
while True:
|
|
await asyncio.sleep(5)
|
|
|
|
a = Actor.remote()
|
|
a.loop_forever.remote()
|
|
pid = ray.get(a.getpid.remote())
|
|
with pytest.raises(ray.exceptions.RayActorError):
|
|
ray.get(a.exit.remote())
|
|
wait_for_pid_to_exit(pid)
|
|
|
|
|
|
def test_async_callback(ray_start_regular_shared):
|
|
global_set = set()
|
|
|
|
ref = ray.put(None)
|
|
ref._on_completed(lambda _: global_set.add("completed-1"))
|
|
wait_for_condition(lambda: "completed-1" in global_set)
|
|
|
|
signal = SignalActor.remote()
|
|
|
|
@ray.remote
|
|
def wait():
|
|
ray.get(signal.wait.remote())
|
|
|
|
ref = wait.remote()
|
|
ref._on_completed(lambda _: global_set.add("completed-2"))
|
|
assert "completed-2" not in global_set
|
|
signal.send.remote()
|
|
wait_for_condition(lambda: "completed-2" in global_set)
|
|
|
|
|
|
@pytest.mark.parametrize("raise_in_callback", [False, True])
|
|
@pytest.mark.skipif(
|
|
client_mode_should_convert(), reason="Different ref counting in Ray client."
|
|
)
|
|
def test_on_completed_callback_refcount(ray_start_regular_shared, raise_in_callback):
|
|
"""Check that the _on_completed callback is ref counted properly."""
|
|
signal = SignalActor.remote()
|
|
|
|
def callback(result):
|
|
if raise_in_callback:
|
|
raise Exception("ruh-roh")
|
|
|
|
@ray.remote
|
|
def wait():
|
|
ray.get(signal.wait.remote())
|
|
|
|
ref = wait.remote()
|
|
|
|
initial_refcount = sys.getrefcount(callback)
|
|
ref._on_completed(callback)
|
|
|
|
# Python ref count should be incremented to avoid the callback being GC'd while the
|
|
# C++ core worker holds a ref to it.
|
|
assert sys.getrefcount(callback) > initial_refcount
|
|
|
|
# Trigger the task to finish so the callback should execute.
|
|
ray.get(signal.send.remote())
|
|
|
|
# Now the refcount should drop back down to the initial count.
|
|
wait_for_condition(lambda: sys.getrefcount(callback) == initial_refcount)
|
|
|
|
|
|
def test_async_function_errored(ray_start_regular_shared):
|
|
with pytest.raises(ValueError):
|
|
|
|
@ray.remote
|
|
async def f():
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_obj_unhandled_errors(ray_start_regular_shared):
|
|
@ray.remote
|
|
def f():
|
|
raise ValueError()
|
|
|
|
num_exceptions = 0
|
|
|
|
def interceptor(e):
|
|
nonlocal num_exceptions
|
|
num_exceptions += 1
|
|
|
|
# Test we report unhandled exceptions.
|
|
ray._private.worker._unhandled_error_handler = interceptor
|
|
x1 = f.remote()
|
|
# NOTE: Unhandled exception is from waiting for the value of x1's ObjectID
|
|
# in x1's destructor, and receiving an exception from f() instead.
|
|
del x1
|
|
wait_for_condition(lambda: num_exceptions == 1)
|
|
|
|
# Test we don't report handled exceptions.
|
|
x1 = f.remote()
|
|
with pytest.raises(ray.exceptions.RayError):
|
|
await x1
|
|
del x1
|
|
await asyncio.sleep(1)
|
|
assert num_exceptions == 1, num_exceptions
|
|
|
|
|
|
# This case tests that the asyncio actor shouldn't create thread
|
|
# pool with max_concurrency threads. Otherwise it will allocate
|
|
# too many resources for threads to lead worker crash.
|
|
def test_asyncio_actor_with_large_concurrency(ray_start_regular_shared):
|
|
@ray.remote
|
|
class Actor:
|
|
def sync_thread_id(self):
|
|
time.sleep(2)
|
|
return threading.current_thread().ident
|
|
|
|
async def async_thread_id(self):
|
|
time.sleep(2)
|
|
return threading.current_thread().ident
|
|
|
|
a = Actor.options(max_concurrency=100000).remote()
|
|
sync_id, async_id = ray.get([a.sync_thread_id.remote(), a.async_thread_id.remote()])
|
|
assert sync_id == async_id
|
|
|
|
|
|
def test_asyncio_actor_shutdown_when_non_async_method_mixed(ray_start_regular_shared):
|
|
# Regression test for: https://github.com/ray-project/ray/issues/32376
|
|
# Ensure the core worker doesn't crash when exit_actor is used while mixing async
|
|
# and sync actor tasks.
|
|
@ray.remote
|
|
class A:
|
|
def __init__(self, *, exit_after: int):
|
|
self._remaining = exit_after
|
|
self._event = asyncio.Event()
|
|
|
|
async def wait_then_exit(self):
|
|
await self._event.wait()
|
|
ray.actor.exit_actor()
|
|
|
|
def ping(self):
|
|
self._remaining -= 1
|
|
if self._remaining == 0:
|
|
self._event.set()
|
|
|
|
# Exit after 1/2 of the ping tasks have executed to ensure interleaving.
|
|
a = A.remote(exit_after=500)
|
|
exit_ref = a.wait_then_exit.remote()
|
|
ping_refs = [a.ping.remote() for _ in range(1000)]
|
|
|
|
with pytest.raises(
|
|
ray.exceptions.RayActorError,
|
|
match="INTENDED_USER_EXIT",
|
|
):
|
|
ray.get([exit_ref] + ping_refs)
|
|
|
|
|
|
def test_asyncio_actor_argument_collision(ray_start_regular_shared):
|
|
"""Regression test for https://github.com/ray-project/ray/issues/41272."""
|
|
|
|
@ray.remote
|
|
class A:
|
|
async def hi_async(self, task_id: str, specified_cgname: str):
|
|
return f"Hi from async: {task_id}! cgname: {specified_cgname}."
|
|
|
|
def hi_sync(self, task_id: str, *, specified_cgname: str):
|
|
return f"Hi from sync: {task_id}! cgname: {specified_cgname}."
|
|
|
|
a = A.remote()
|
|
assert (
|
|
ray.get(a.hi_async.remote(task_id="TEST", specified_cgname="test2"))
|
|
== "Hi from async: TEST! cgname: test2."
|
|
)
|
|
assert (
|
|
ray.get(a.hi_sync.remote(task_id="TEST", specified_cgname="test2"))
|
|
== "Hi from sync: TEST! cgname: test2."
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|