wehub-resource-sync
304 lines
7.5 KiB
Python
304 lines
7.5 KiB
Python
import asyncio
|
|
import sys
|
|
import time
|
|
from collections import Counter
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._raylet import ObjectRefGenerator
|
|
from ray.exceptions import TaskCancelledError
|
|
|
|
|
|
def test_threaded_actor_generator(shutdown_only):
|
|
ray.init()
|
|
|
|
@ray.remote(max_concurrency=10)
|
|
class Actor:
|
|
def f(self):
|
|
for i in range(30):
|
|
time.sleep(0.1)
|
|
yield np.ones(1024 * 1024) * i
|
|
|
|
@ray.remote(max_concurrency=20)
|
|
class AsyncActor:
|
|
async def f(self):
|
|
for i in range(30):
|
|
await asyncio.sleep(0.1)
|
|
yield np.ones(1024 * 1024) * i
|
|
|
|
async def main():
|
|
a = Actor.remote()
|
|
asy = AsyncActor.remote()
|
|
|
|
async def run():
|
|
i = 0
|
|
async for ref in a.f.remote():
|
|
val = ray.get(ref)
|
|
print(val)
|
|
print(ref)
|
|
assert np.array_equal(val, np.ones(1024 * 1024) * i)
|
|
i += 1
|
|
del ref
|
|
|
|
async def run2():
|
|
i = 0
|
|
async for ref in asy.f.remote():
|
|
val = await ref
|
|
print(ref)
|
|
print(val)
|
|
assert np.array_equal(val, np.ones(1024 * 1024) * i), ref
|
|
i += 1
|
|
del ref
|
|
|
|
coroutines = [run() for _ in range(10)]
|
|
coroutines = [run2() for _ in range(20)]
|
|
|
|
await asyncio.gather(*coroutines)
|
|
|
|
asyncio.run(main())
|
|
|
|
|
|
def test_generator_dist_gather(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=0, object_store_memory=1 * 1024 * 1024 * 1024)
|
|
ray.init()
|
|
cluster.add_node(num_cpus=1)
|
|
cluster.add_node(num_cpus=1)
|
|
cluster.add_node(num_cpus=1)
|
|
cluster.add_node(num_cpus=1)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def __init__(self, child=None):
|
|
self.child = child
|
|
|
|
def get_data(self):
|
|
for _ in range(10):
|
|
time.sleep(0.1)
|
|
yield np.ones(5 * 1024 * 1024)
|
|
|
|
async def all_gather():
|
|
actor = Actor.remote()
|
|
async for ref in actor.get_data.remote():
|
|
val = await ref
|
|
assert np.array_equal(np.ones(5 * 1024 * 1024), val)
|
|
del ref
|
|
|
|
async def main():
|
|
await asyncio.gather(all_gather(), all_gather(), all_gather(), all_gather())
|
|
|
|
asyncio.run(main())
|
|
summary = ray._private.internal_api.memory_summary(stats_only=True)
|
|
print(summary)
|
|
|
|
|
|
def test_generator_wait(shutdown_only):
|
|
"""
|
|
Make sure the generator works with ray.wait.
|
|
"""
|
|
ray.init(num_cpus=8)
|
|
|
|
@ray.remote
|
|
def f(sleep_time):
|
|
for i in range(2):
|
|
time.sleep(sleep_time)
|
|
yield i
|
|
|
|
@ray.remote
|
|
def g(sleep_time):
|
|
time.sleep(sleep_time)
|
|
return 10
|
|
|
|
gen = f.remote(1)
|
|
|
|
"""
|
|
Test basic cases.
|
|
"""
|
|
for expected_rval in [0, 1]:
|
|
s = time.time()
|
|
r, ur = ray.wait([gen], num_returns=1)
|
|
print(time.time() - s)
|
|
assert len(r) == 1
|
|
assert ray.get(next(r[0])) == expected_rval
|
|
assert len(ur) == 0
|
|
|
|
# Should raise a stop iteration.
|
|
for _ in range(3):
|
|
s = time.time()
|
|
r, ur = ray.wait([gen], num_returns=1)
|
|
print(time.time() - s)
|
|
assert len(r) == 1
|
|
with pytest.raises(StopIteration):
|
|
assert next(r[0]) == 0
|
|
assert len(ur) == 0
|
|
|
|
gen = f.remote(0)
|
|
# Wait until the generator task finishes
|
|
ray.get(gen._generator_ref)
|
|
for i in range(2):
|
|
r, ur = ray.wait([gen], timeout=0)
|
|
assert len(r) == 1
|
|
assert len(ur) == 0
|
|
assert ray.get(next(r[0])) == i
|
|
|
|
"""
|
|
Test the case ref is mixed with regular object ref.
|
|
"""
|
|
gen = f.remote(0)
|
|
ref = g.remote(3)
|
|
ready, unready = [], [gen, ref]
|
|
result_set = set()
|
|
while unready:
|
|
ready, unready = ray.wait(unready)
|
|
print(ready, unready)
|
|
assert len(ready) == 1
|
|
for r in ready:
|
|
if isinstance(r, ObjectRefGenerator):
|
|
try:
|
|
ref = next(r)
|
|
print(ref)
|
|
print(ray.get(ref))
|
|
result_set.add(ray.get(ref))
|
|
except StopIteration:
|
|
pass
|
|
else:
|
|
unready.append(r)
|
|
else:
|
|
result_set.add(ray.get(r))
|
|
|
|
assert result_set == {0, 1, 10}
|
|
|
|
"""
|
|
Test timeout.
|
|
"""
|
|
gen = f.remote(3)
|
|
ref = g.remote(1)
|
|
ready, unready = ray.wait([gen, ref], timeout=2)
|
|
assert len(ready) == 1
|
|
assert len(unready) == 1
|
|
|
|
"""
|
|
Test num_returns
|
|
"""
|
|
gen = f.remote(1)
|
|
ref = g.remote(1)
|
|
ready, unready = ray.wait([ref, gen], num_returns=2)
|
|
assert len(ready) == 2
|
|
assert len(unready) == 0
|
|
|
|
|
|
@pytest.mark.parametrize("backpressure", [True, False])
|
|
def test_generator_wait_e2e(shutdown_only, backpressure):
|
|
ray.init(num_cpus=8)
|
|
|
|
if backpressure:
|
|
threshold = 1
|
|
else:
|
|
threshold = -1
|
|
|
|
@ray.remote
|
|
def f(sleep_time):
|
|
for i in range(2):
|
|
time.sleep(sleep_time)
|
|
yield i
|
|
|
|
@ray.remote
|
|
def g(sleep_time):
|
|
time.sleep(sleep_time)
|
|
return 10
|
|
|
|
gen = [
|
|
f.options(
|
|
_generator_backpressure_num_objects=threshold,
|
|
).remote(1)
|
|
for _ in range(4)
|
|
]
|
|
ref = [g.remote(2) for _ in range(4)]
|
|
ready, unready = [], [*gen, *ref]
|
|
result = []
|
|
start = time.time()
|
|
while unready:
|
|
ready, unready = ray.wait(unready, num_returns=len(unready), timeout=0.1)
|
|
for r in ready:
|
|
if isinstance(r, ObjectRefGenerator):
|
|
try:
|
|
ref = next(r)
|
|
result.append(ray.get(ref))
|
|
except StopIteration:
|
|
pass
|
|
else:
|
|
unready.append(r)
|
|
else:
|
|
result.append(ray.get(r))
|
|
elapsed = time.time() - start
|
|
assert elapsed < 4
|
|
assert 2 < elapsed
|
|
|
|
assert len(result) == 12
|
|
result = Counter(result)
|
|
assert result[0] == 4
|
|
assert result[1] == 4
|
|
assert result[10] == 4
|
|
|
|
|
|
def test_completed_next_ready_is_finished(shutdown_only):
|
|
@ray.remote
|
|
def f():
|
|
for _ in range(3):
|
|
time.sleep(1)
|
|
yield 1
|
|
|
|
gen = f.remote()
|
|
assert not gen.is_finished()
|
|
assert not gen.next_ready()
|
|
r, _ = ray.wait([gen])
|
|
gen = r[0]
|
|
assert gen.next_ready()
|
|
_, ur = ray.wait([gen.completed()], timeout=0)
|
|
assert len(ur) == 1
|
|
|
|
# Consume object refs
|
|
next(gen)
|
|
assert not gen.is_finished()
|
|
_, ur = ray.wait([gen.completed()], timeout=0)
|
|
assert len(ur) == 1
|
|
|
|
next(gen)
|
|
assert not gen.is_finished()
|
|
_, ur = ray.wait([gen.completed()], timeout=0)
|
|
assert len(ur) == 1
|
|
|
|
next(gen)
|
|
with pytest.raises(StopIteration):
|
|
next(gen)
|
|
|
|
assert gen.is_finished()
|
|
# Since the next should raise StopIteration,
|
|
# it should be False.
|
|
assert not gen.next_ready()
|
|
r, _ = ray.wait([gen.completed()], timeout=0)
|
|
assert len(r) == 1
|
|
|
|
# Test the failed case.
|
|
gen = f.remote()
|
|
next(gen)
|
|
ray.cancel(gen, force=True)
|
|
r, _ = ray.wait([gen])
|
|
assert len(r) == 1
|
|
# The last exception is not taken yet.
|
|
assert gen.next_ready()
|
|
assert not gen.is_finished()
|
|
with pytest.raises(TaskCancelledError):
|
|
ray.get(gen.completed())
|
|
with pytest.raises(TaskCancelledError):
|
|
ray.get(next(gen))
|
|
assert not gen.next_ready()
|
|
assert gen.is_finished()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
sys.exit(pytest.main(["-sv", __file__]))
|