wehub-resource-sync
363 lines
8.9 KiB
Python
363 lines
8.9 KiB
Python
import pytest
|
|
|
|
import ray
|
|
from ray.dag import (
|
|
PARENT_CLASS_NODE_KEY,
|
|
PREV_CLASS_METHOD_CALL_KEY,
|
|
InputNode,
|
|
MultiOutputNode,
|
|
)
|
|
|
|
|
|
@ray.remote
|
|
class Counter:
|
|
def __init__(self, init_value=0):
|
|
self.i = init_value
|
|
|
|
def inc(self):
|
|
self.i += 1
|
|
|
|
def get(self):
|
|
return self.i
|
|
|
|
|
|
@ray.remote
|
|
class Actor:
|
|
def __init__(self, init_value):
|
|
self.i = init_value
|
|
|
|
def inc(self, x):
|
|
self.i += x
|
|
|
|
def get(self):
|
|
return self.i
|
|
|
|
def echo(self, x):
|
|
return x
|
|
|
|
@ray.method(num_returns=2)
|
|
def return_two(self, x):
|
|
return x, x + 1
|
|
|
|
@ray.method(num_returns=2)
|
|
def inc_and_return_two(self, x):
|
|
self.i += x
|
|
return self.i, self.i + 1
|
|
|
|
@ray.method(num_returns=1)
|
|
def return_two_as_one(self, x):
|
|
return x, x + 1
|
|
|
|
@ray.method(num_returns=2)
|
|
def return_two_from_three(self, x):
|
|
return x, x + 1, x + 2
|
|
|
|
|
|
def test_basic_actor_dag(shared_ray_instance):
|
|
@ray.remote
|
|
def combine(x, y):
|
|
return x + y
|
|
|
|
a1 = Actor.bind(10)
|
|
res = a1.get.bind()
|
|
print(res)
|
|
assert ray.get(res.execute()) == 10
|
|
|
|
a2 = Actor.bind(10)
|
|
a1.inc.bind(2)
|
|
a1.inc.bind(4)
|
|
a2.inc.bind(6)
|
|
dag = combine.bind(a1.get.bind(), a2.get.bind())
|
|
|
|
print(dag)
|
|
assert ray.get(dag.execute()) == 32
|
|
|
|
|
|
def test_class_as_class_constructor_arg(shared_ray_instance):
|
|
@ray.remote
|
|
class OuterActor:
|
|
def __init__(self, inner_actor):
|
|
self.inner_actor = inner_actor
|
|
|
|
def inc(self, x):
|
|
self.inner_actor.inc.remote(x)
|
|
|
|
def get(self):
|
|
return ray.get(self.inner_actor.get.remote())
|
|
|
|
outer = OuterActor.bind(Actor.bind(10))
|
|
outer.inc.bind(2)
|
|
dag = outer.get.bind()
|
|
print(dag)
|
|
assert ray.get(dag.execute()) == 12
|
|
|
|
|
|
def test_class_as_function_constructor_arg(shared_ray_instance):
|
|
@ray.remote
|
|
def f(actor_handle):
|
|
return ray.get(actor_handle.get.remote())
|
|
|
|
dag = f.bind(Actor.bind(10))
|
|
print(dag)
|
|
assert ray.get(dag.execute()) == 10
|
|
|
|
|
|
def test_basic_actor_dag_constructor_options(shared_ray_instance):
|
|
a1 = Actor.bind(10)
|
|
dag = a1.get.bind()
|
|
print(dag)
|
|
assert ray.get(dag.execute()) == 10
|
|
|
|
a1 = Actor.options(name="Actor", namespace="test", max_pending_calls=10).bind(10)
|
|
dag = a1.get.bind()
|
|
print(dag)
|
|
# Ensure execution result is identical with .options() in init()
|
|
assert ray.get(dag.execute()) == 10
|
|
# Ensure options are passed in
|
|
assert a1.get_options().get("name") == "Actor"
|
|
assert a1.get_options().get("namespace") == "test"
|
|
assert a1.get_options().get("max_pending_calls") == 10
|
|
|
|
|
|
def test_actor_method_options(shared_ray_instance):
|
|
a1 = Actor.bind(10)
|
|
dag = a1.get.options(name="actor_method_options").bind()
|
|
print(dag)
|
|
assert ray.get(dag.execute()) == 10
|
|
assert dag.get_options().get("name") == "actor_method_options"
|
|
|
|
|
|
def test_basic_actor_dag_constructor_invalid_options(shared_ray_instance):
|
|
with pytest.raises(
|
|
ValueError, match=r".*quantity of resource num_cpus cannot be negative.*"
|
|
):
|
|
a1 = Actor.options(num_cpus=-1).bind(10)
|
|
invalid_dag = a1.get.bind()
|
|
ray.get(invalid_dag.execute())
|
|
|
|
|
|
def test_actor_options_complicated(shared_ray_instance):
|
|
"""Test a more complicated setup where we apply .options() in both
|
|
constructor and method call with overlapping keys, and ensure end to end
|
|
options correctness.
|
|
"""
|
|
|
|
@ray.remote
|
|
def combine(x, y):
|
|
return x + y
|
|
|
|
a1 = Actor.options(name="a1_v0").bind(10)
|
|
res = a1.get.options(name="v1").bind()
|
|
print(res)
|
|
assert ray.get(res.execute()) == 10
|
|
assert a1.get_options().get("name") == "a1_v0"
|
|
assert res.get_options().get("name") == "v1"
|
|
|
|
a1 = Actor.options(name="a1_v1").bind(10) # Cannot
|
|
a2 = Actor.options(name="a2_v0").bind(10)
|
|
a1.inc.options(name="v1").bind(2)
|
|
a1.inc.options(name="v2").bind(4)
|
|
a2.inc.options(name="v3").bind(6)
|
|
dag = combine.options(name="v4").bind(a1.get.bind(), a2.get.bind())
|
|
|
|
print(dag)
|
|
assert ray.get(dag.execute()) == 32
|
|
test_a1 = dag.get_args()[0] # call graph for a1.get.bind()
|
|
test_a2 = dag.get_args()[1] # call graph for a2.get.bind()
|
|
assert test_a2.get_options() == {} # No .options() at outer call
|
|
# refer to a2 constructor .options() call
|
|
assert (
|
|
test_a2.get_other_args_to_resolve()[PARENT_CLASS_NODE_KEY]
|
|
.get_options()
|
|
.get("name")
|
|
== "a2_v0"
|
|
)
|
|
# refer to actor method a2.inc.options() call
|
|
assert (
|
|
test_a2.get_other_args_to_resolve()[PREV_CLASS_METHOD_CALL_KEY]
|
|
.get_options()
|
|
.get("name")
|
|
== "v3"
|
|
)
|
|
# refer to a1 constructor .options() call
|
|
assert (
|
|
test_a1.get_other_args_to_resolve()[PARENT_CLASS_NODE_KEY]
|
|
.get_options()
|
|
.get("name")
|
|
== "a1_v1"
|
|
)
|
|
# refer to latest actor method a1.inc.options() call
|
|
assert (
|
|
test_a1.get_other_args_to_resolve()[PREV_CLASS_METHOD_CALL_KEY]
|
|
.get_options()
|
|
.get("name")
|
|
== "v2"
|
|
)
|
|
# refer to first bound actor method a1.inc.options() call
|
|
assert (
|
|
test_a1.get_other_args_to_resolve()[PREV_CLASS_METHOD_CALL_KEY]
|
|
.get_other_args_to_resolve()[PREV_CLASS_METHOD_CALL_KEY]
|
|
.get_options()
|
|
.get("name")
|
|
== "v1"
|
|
)
|
|
|
|
|
|
def test_pass_actor_handle(shared_ray_instance):
|
|
@ray.remote
|
|
class Actor:
|
|
def ping(self):
|
|
return "hello"
|
|
|
|
@ray.remote
|
|
def caller(handle):
|
|
assert isinstance(handle, ray.actor.ActorHandle), handle
|
|
return ray.get(handle.ping.remote())
|
|
|
|
a1 = Actor.bind()
|
|
dag = caller.bind(a1)
|
|
print(dag)
|
|
assert ray.get(dag.execute()) == "hello"
|
|
|
|
|
|
def test_dynamic_pipeline(shared_ray_instance):
|
|
@ray.remote
|
|
class Model:
|
|
def __init__(self, arg):
|
|
self.arg = arg
|
|
|
|
def forward(self, x):
|
|
return self.arg + str(x)
|
|
|
|
@ray.remote
|
|
class ModelSelection:
|
|
def is_even(self, x):
|
|
return x % 2 == 0
|
|
|
|
@ray.remote
|
|
def pipeline(x, m1, m2, selection):
|
|
sel = selection.is_even.remote(x)
|
|
if ray.get(sel):
|
|
result = m1.forward.remote(x)
|
|
else:
|
|
result = m2.forward.remote(x)
|
|
return ray.get(result)
|
|
|
|
m1 = Model.bind("Even: ")
|
|
m2 = Model.bind("Odd: ")
|
|
selection = ModelSelection.bind()
|
|
|
|
even_input = pipeline.bind(20, m1, m2, selection)
|
|
print(even_input)
|
|
assert ray.get(even_input.execute()) == "Even: 20"
|
|
|
|
odd_input = pipeline.bind(21, m1, m2, selection)
|
|
print(odd_input)
|
|
assert ray.get(odd_input.execute()) == "Odd: 21"
|
|
|
|
|
|
def test_unsupported_bind():
|
|
@ray.remote
|
|
class Actor:
|
|
def ping(self):
|
|
return "hello"
|
|
|
|
with pytest.raises(
|
|
AttributeError,
|
|
match=r"\.bind\(\) cannot be used again on",
|
|
):
|
|
actor = Actor.bind()
|
|
_ = actor.bind()
|
|
|
|
with pytest.raises(
|
|
AttributeError,
|
|
match=r"\.remote\(\) cannot be used on ClassMethodNodes",
|
|
):
|
|
actor = Actor.bind()
|
|
_ = actor.ping.remote()
|
|
|
|
|
|
def test_unsupported_remote():
|
|
@ray.remote
|
|
class Actor:
|
|
def ping(self):
|
|
return "hello"
|
|
|
|
with pytest.raises(AttributeError, match="'Actor' has no attribute 'remote'"):
|
|
_ = Actor.bind().remote()
|
|
|
|
@ray.remote
|
|
def func():
|
|
return 1
|
|
|
|
with pytest.raises(AttributeError, match=r"\.remote\(\) cannot be used on"):
|
|
_ = func.bind().remote()
|
|
|
|
|
|
def test_two_returns_first():
|
|
a = Actor.remote(0)
|
|
with InputNode() as i:
|
|
o1, o2 = a.return_two.bind(i)
|
|
dag = o1
|
|
|
|
for _ in range(3):
|
|
res = ray.get(dag.execute(1))
|
|
assert res == 1
|
|
|
|
|
|
def test_two_returns_second():
|
|
a = Actor.remote(0)
|
|
with InputNode() as i:
|
|
o1, o2 = a.return_two.bind(i)
|
|
dag = o2
|
|
|
|
for _ in range(3):
|
|
res = ray.get(dag.execute(1))
|
|
assert res == 2
|
|
|
|
|
|
def test_two_returns_one_reader_multi_times():
|
|
a = Actor.remote(0)
|
|
b = Actor.remote(0)
|
|
with InputNode() as i:
|
|
o1, o2 = a.return_two.bind(i)
|
|
o3 = b.echo.bind(o1)
|
|
o4 = b.echo.bind(o2)
|
|
dag = MultiOutputNode([o3, o4])
|
|
|
|
for _ in range(3):
|
|
res = ray.get(dag.execute(1))
|
|
assert res == [1, 2]
|
|
|
|
|
|
def test_two_returns_two_readers():
|
|
a = Actor.remote(0)
|
|
b = Actor.remote(0)
|
|
c = Actor.remote(0)
|
|
with InputNode() as i:
|
|
o1, o2 = a.return_two.bind(i)
|
|
o3 = b.echo.bind(o1)
|
|
o4 = c.echo.bind(o2)
|
|
dag = MultiOutputNode([o3, o4])
|
|
|
|
for _ in range(3):
|
|
res = ray.get(dag.execute(1))
|
|
assert res == [1, 2]
|
|
|
|
|
|
def test_inc_two_returns():
|
|
a = Actor.remote(0)
|
|
with InputNode() as i:
|
|
o1, o2 = a.inc_and_return_two.bind(i)
|
|
dag = MultiOutputNode([o1, o2])
|
|
|
|
for i in range(3):
|
|
res = ray.get(dag.execute(1))
|
|
assert res == [i + 1, i + 2]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|