wehub-resource-sync
279 lines
6.3 KiB
Python
279 lines
6.3 KiB
Python
import asyncio
|
|
import sys
|
|
import time
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.util import ActorPool
|
|
|
|
|
|
@pytest.fixture
|
|
def init():
|
|
ray.init(num_cpus=4)
|
|
yield
|
|
ray.shutdown()
|
|
|
|
|
|
def test_get_next(init):
|
|
@ray.remote
|
|
class MyActor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def f(self, x):
|
|
return x + 1
|
|
|
|
def double(self, x):
|
|
return 2 * x
|
|
|
|
actors = [MyActor.remote() for _ in range(4)]
|
|
pool = ActorPool(actors)
|
|
for i in range(5):
|
|
pool.submit(lambda a, v: a.f.remote(v), i)
|
|
assert pool.get_next() == i + 1
|
|
|
|
|
|
def test_get_next_unordered(init):
|
|
@ray.remote
|
|
class MyActor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def f(self, x):
|
|
return x + 1
|
|
|
|
def double(self, x):
|
|
return 2 * x
|
|
|
|
actors = [MyActor.remote() for _ in range(4)]
|
|
pool = ActorPool(actors)
|
|
|
|
total = []
|
|
|
|
for i in range(5):
|
|
pool.submit(lambda a, v: a.f.remote(v), i)
|
|
while pool.has_next():
|
|
total += [pool.get_next_unordered()]
|
|
|
|
assert all(elem in [1, 2, 3, 4, 5] for elem in total)
|
|
|
|
|
|
def test_map(init):
|
|
@ray.remote
|
|
class MyActor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def f(self, x):
|
|
return x + 1
|
|
|
|
def double(self, x):
|
|
return 2 * x
|
|
|
|
actors = [MyActor.remote() for _ in range(4)]
|
|
pool = ActorPool(actors)
|
|
|
|
index = 0
|
|
for v in pool.map(lambda a, v: a.double.remote(v), range(5)):
|
|
assert v == 2 * index
|
|
index += 1
|
|
|
|
|
|
def test_map_eager(init):
|
|
"""Verify that submit is called eagerly when map is called.
|
|
|
|
If the results are directly yielded, then the submit calls are not
|
|
executed until the results are consumed.
|
|
"""
|
|
|
|
@ray.remote
|
|
class MyActor:
|
|
def f(self, x):
|
|
pass
|
|
|
|
actor = MyActor.remote()
|
|
pool = ActorPool([actor])
|
|
pool.submit = MagicMock()
|
|
|
|
pool.map(lambda a, v: a.f.remote(v), range(1))
|
|
pool.submit.assert_called()
|
|
|
|
|
|
def test_map_unordered(init):
|
|
@ray.remote
|
|
class MyActor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def f(self, x):
|
|
return x + 1
|
|
|
|
def double(self, x):
|
|
return 2 * x
|
|
|
|
actors = [MyActor.remote() for _ in range(4)]
|
|
pool = ActorPool(actors)
|
|
|
|
total = []
|
|
for v in pool.map(lambda a, v: a.double.remote(v), range(5)):
|
|
total += [v]
|
|
|
|
assert all(elem in [0, 2, 4, 6, 8] for elem in total)
|
|
|
|
|
|
def test_map_gh23107(init):
|
|
sleep_time = 40
|
|
|
|
# Reference - https://github.com/ray-project/ray/issues/23107
|
|
@ray.remote
|
|
class DummyActor:
|
|
async def identity(self, s):
|
|
if s == 6:
|
|
await asyncio.sleep(sleep_time)
|
|
return s, time.time()
|
|
|
|
def func(a, v):
|
|
return a.identity.remote(v)
|
|
|
|
map_values = [1, 2, 3, 4, 5]
|
|
|
|
pool_map = ActorPool([DummyActor.remote() for i in range(2)])
|
|
pool_map.submit(func, 6)
|
|
start_time = time.time()
|
|
gen = pool_map.map(func, map_values)
|
|
assert all(elem[0] in [1, 2, 3, 4, 5] for elem in list(gen))
|
|
assert all(
|
|
abs(elem[1] - start_time) < sleep_time in [1, 2, 3, 4, 5] for elem in list(gen)
|
|
)
|
|
|
|
pool_map_unordered = ActorPool([DummyActor.remote() for i in range(2)])
|
|
pool_map_unordered.submit(func, 6)
|
|
start_time = time.time()
|
|
gen = pool_map_unordered.map_unordered(func, map_values)
|
|
assert all(elem[0] in [1, 2, 3, 4, 5] for elem in list(gen))
|
|
assert all(
|
|
abs(elem[1] - start_time) < sleep_time in [1, 2, 3, 4, 5] for elem in list(gen)
|
|
)
|
|
|
|
|
|
def test_get_next_timeout(init):
|
|
@ray.remote
|
|
class MyActor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def f(self, x):
|
|
while True:
|
|
x = x + 1
|
|
time.sleep(1)
|
|
return None
|
|
|
|
def double(self, x):
|
|
return 2 * x
|
|
|
|
actors = [MyActor.remote() for _ in range(4)]
|
|
pool = ActorPool(actors)
|
|
pool.submit(lambda a, v: a.f.remote(v), 0)
|
|
with pytest.raises(TimeoutError):
|
|
pool.get_next_unordered(timeout=0.1)
|
|
|
|
|
|
def test_get_next_unordered_timeout(init):
|
|
@ray.remote
|
|
class MyActor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def f(self, x):
|
|
while True:
|
|
x + 1
|
|
time.sleep(1)
|
|
return
|
|
|
|
def double(self, x):
|
|
return 2 * x
|
|
|
|
actors = [MyActor.remote() for _ in range(4)]
|
|
pool = ActorPool(actors)
|
|
|
|
pool.submit(lambda a, v: a.f.remote(v), 0)
|
|
with pytest.raises(TimeoutError):
|
|
pool.get_next_unordered(timeout=0.1)
|
|
|
|
|
|
def test_multiple_returns(init):
|
|
@ray.remote
|
|
class Foo(object):
|
|
@ray.method(num_returns=2)
|
|
def bar(self):
|
|
return 1, 2
|
|
|
|
pool = ActorPool([Foo.remote() for _ in range(2)])
|
|
for _ in range(4):
|
|
pool.submit(lambda a, v: a.bar.remote(), None)
|
|
|
|
while pool.has_next():
|
|
assert pool.get_next(timeout=None) == [1, 2]
|
|
|
|
|
|
def test_pop_idle(init):
|
|
@ray.remote
|
|
class MyActor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def f(self, x):
|
|
return x + 1
|
|
|
|
def double(self, x):
|
|
return 2 * x
|
|
|
|
actors = [MyActor.remote()]
|
|
pool = ActorPool(actors)
|
|
|
|
pool.submit(lambda a, v: a.double.remote(v), 1)
|
|
assert pool.pop_idle() is None
|
|
assert pool.has_free() is False # actor is busy
|
|
assert pool.get_next() == 2
|
|
assert pool.has_free()
|
|
pool.pop_idle() # removes actor from pool
|
|
assert pool.has_free() is False # no more actors in pool
|
|
|
|
|
|
def test_push(init):
|
|
@ray.remote
|
|
class MyActor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def f(self, x):
|
|
return x + 1
|
|
|
|
def double(self, x):
|
|
return 2 * x
|
|
|
|
a1, a2, a3 = MyActor.remote(), MyActor.remote(), MyActor.remote()
|
|
pool = ActorPool([a1])
|
|
|
|
pool.submit(lambda a, v: a.double.remote(v), 1)
|
|
assert pool.has_free() is False # actor is busy
|
|
with pytest.raises(ValueError):
|
|
pool.push(a1)
|
|
pool.push(a2)
|
|
assert pool.has_free() # a2 is available
|
|
|
|
pool.submit(lambda a, v: a.double.remote(v), 1)
|
|
pool.submit(lambda a, v: a.double.remote(v), 1)
|
|
assert pool.has_free() is False
|
|
assert len(pool._pending_submits) == 1
|
|
pool.push(a3)
|
|
assert pool.has_free() is False # a3 is used for pending submit
|
|
assert len(pool._pending_submits) == 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
sys.exit(pytest.main(["-sv", __file__]))
|