def test_no_args_to_input_node(shared_ray_instance):
@ray.remote
def f(input):
return input
with pytest.raises(ValueError,
match="InputNode should not take any args or kwargs"):
with InputNode(0) as dag_input:
f.bind(dag_input)
with pytest.raises(
ValueError,
match="InputNode should not take any args or kwargs",
):
with InputNode(key=1) as dag_input:
f.bind(dag_input)
def get_shared_deployment_handle_dag():
with InputNode() as dag_input:
m = Model.bind(2)
combine = Combine.bind(m, m2=m)
ray_dag = combine.__call__.bind(dag_input)
return ray_dag, dag_input
def test_class_method_input(shared_ray_instance):
@ray.remote
class Model:
def __init__(self, weight: int):
self.weight = weight
def forward(self, input: "RayHandleLike"):
return self.weight * input
@ray.remote
class FeatureProcessor:
def __init__(self, scale):
self.scale = scale
def process(self, input: int):
return input * self.scale
with InputNode() as dag_input:
preprocess = FeatureProcessor.bind(0.5)
feature = preprocess.process.bind(dag_input)
model = Model.bind(4)
dag = model.forward.bind(feature)
# 2 * 0.5 * 4
assert ray.get(dag.execute(2)) == 4
# 6 * 0.5 * 4
assert ray.get(dag.execute(6)) == 12
def test_func_dag(shared_ray_instance):
@ray.remote
def a(user_input):
return user_input
@ray.remote
def b(x):
return x * 2
@ray.remote
def c(x):
return x + 1
@ray.remote
def d(x, y):
return x + y
with InputNode() as dag_input:
a_ref = a.bind(dag_input)
b_ref = b.bind(a_ref)
c_ref = c.bind(a_ref)
d_ref = d.bind(b_ref, c_ref)
d1_ref = d.bind(d_ref, d_ref)
d2_ref = d.bind(d1_ref, d_ref)
dag = d.bind(d2_ref, d_ref)
# [(2*2 + 2+1) + (2*2 + 2+1)] + [(2*2 + 2+1) + (2*2 + 2+1)]
assert ray.get(dag.execute(2)) == 28
# [(3*2 + 3+1) + (3*2 + 3+1)] + [(3*2 + 3+1) + (3*2 + 3+1)]
assert ray.get(dag.execute(3)) == 40
def test_func_class_mixed_input(shared_ray_instance):
"""
Test both class method and function are used as input in the
same dag.
"""
@ray.remote
class Model:
def __init__(self, weight: int):
self.weight = weight
def forward(self, input: int):
return self.weight * input
@ray.remote
def model_func(input: int):
return input * 2
@ray.remote
def combine(m1: "RayHandleLike", m2: "RayHandleLike"):
return m1 + m2
with InputNode() as dag_input:
m1 = Model.bind(3)
m1_output = m1.forward.bind(dag_input)
m2_output = model_func.bind(dag_input)
dag = combine.bind(m1_output, m2_output)
# 2*3 + 2*2
assert ray.get(dag.execute(2)) == 10
# 3*3 + 3*2
assert ray.get(dag.execute(3)) == 15
def test_multi_class_method_input(shared_ray_instance):
"""
Test a multiple class methods can all be used as inputs in a dag.
"""
@ray.remote
class Model:
def __init__(self, weight: int):
self.weight = weight
def forward(self, input: int):
return self.weight * input
@ray.remote
def combine(m1: "RayHandleLike", m2: "RayHandleLike"):
return m1 + m2
with InputNode() as dag_input:
m1 = Model.bind(2)
m2 = Model.bind(3)
m1_output = m1.forward.bind(dag_input)
m2_output = m2.forward.bind(dag_input)
dag = combine.bind(m1_output, m2_output)
# 1*2 + 1*3
assert ray.get(dag.execute(1)) == 5
# 2*2 + 2*3
assert ray.get(dag.execute(2)) == 10
def test_autoscaling_0_replica(serve_instance):
autoscaling_config = {
"metrics_interval_s": 0.1,
"min_replicas": 0,
"max_replicas": 2,
"look_back_period_s": 0.4,
"downscale_delay_s": 0,
"upscale_delay_s": 0,
}
@serve.deployment(
autoscaling_config=autoscaling_config, )
class Model:
def __init__(self, weight):
self.weight = weight
def forward(self, input):
return input + self.weight
with InputNode() as user_input:
model = Model.bind(1)
output = model.forward.bind(user_input)
serve_dag = DAGDriver.options(
route_prefix="/my-dag",
autoscaling_config=autoscaling_config,
).bind(output)
dag_handle = serve.run(serve_dag)
assert 2 == ray.get(dag_handle.predict.remote(1))
def get_multi_instantiation_class_deployment_in_init_args_dag():
with InputNode() as dag_input:
m1 = Model.bind(2)
m2 = Model.bind(3)
combine = Combine.bind(m1, m2=m2)
ray_dag = combine.__call__.bind(dag_input)
return ray_dag, dag_input
def get_multi_instantiation_class_nested_deployment_arg_dag():
with InputNode() as dag_input:
m1 = Model.bind(2)
m2 = Model.bind(3)
combine = Combine.bind(m1, m2={NESTED_HANDLE_KEY: m2}, m2_nested=True)
ray_dag = combine.__call__.bind(dag_input)
return ray_dag, dag_input
def test_ensure_input_node_singleton(shared_ray_instance):
@ray.remote
def f(input):
return input
@ray.remote
def combine(a, b):
return a + b
with InputNode() as input_1:
a = f.bind(input_1)
with InputNode() as input_2:
b = f.bind(input_2)
dag = combine.bind(a, b)
with pytest.raises(AssertionError,
match="Each DAG should only have one unique InputNode"):
_ = ray.get(dag.execute(2))
def get_func_class_with_class_method_dag():
with InputNode() as dag_input:
m1 = Model.bind(1)
m2 = Model.bind(2)
m1_output = m1.forward.bind(dag_input[0])
m2_output = m2.forward.bind(dag_input[1])
ray_dag = combine.bind(m1_output,
m2_output,
kwargs_output=dag_input[2])
return ray_dag, dag_input
def test_multi_instantiation_class_nested_deployment_arg(serve_instance):
with InputNode() as dag_input:
m1 = Model.bind(2)
m2 = Model.bind(3)
combine = Combine.bind(m1, m2={NESTED_HANDLE_KEY: m2}, m2_nested=True)
ray_dag = combine.__call__.bind(dag_input)
(
serve_root_dag,
deserialized_serve_root_dag_node,
) = _test_deployment_json_serde_helper(ray_dag,
input=1,
expected_num_deployments=3)
assert ray.get(serve_root_dag.execute(1)) == ray.get(
deserialized_serve_root_dag_node.execute(1))
def test_ensure_in_context_manager(shared_ray_instance):
# No enforcement on creation given __enter__ executes after __init__
input = InputNode()
with pytest.raises(
AssertionError,
match=(
"InputNode is a singleton instance that should be only used "
"in context manager"),
):
input.execute()
@ray.remote
def f(input):
return input
# No enforcement on creation given __enter__ executes after __init__
dag = f.bind(InputNode())
with pytest.raises(
AssertionError,
match=(
"InputNode is a singleton instance that should be only used "
"in context manager"),
):
dag.execute()
def test_simple_func(shared_ray_instance):
@ray.remote
def a(input: str):
return f"{input} -> a"
@ray.remote
def b(a: "RayHandleLike"):
# At runtime, a is replaced with execution result of a.
return f"{a} -> b"
# input -> a - > b -> ouput
with InputNode() as dag_input:
a_node = a.bind(dag_input)
dag = b.bind(a_node)
assert ray.get(dag.execute("input")) == "input -> a -> b"
assert ray.get(dag.execute("test")) == "test -> a -> b"
def test_nested_deployment_node_json_serde(serve_instance):
with InputNode() as dag_input:
m1 = Model.bind(2)
m2 = Model.bind(3)
m1_output = m1.forward.bind(dag_input)
m2_output = m2.forward.bind(dag_input)
ray_dag = combine.bind(m1_output, m2_output)
(
serve_root_dag,
deserialized_serve_root_dag_node,
) = _test_deployment_json_serde_helper(ray_dag,
input=1,
expected_num_deployments=2)
assert ray.get(serve_root_dag.execute(1)) == ray.get(
deserialized_serve_root_dag_node.execute(1))
def test_invalid_input_node_as_class_constructor(shared_ray_instance):
@ray.remote
class Actor:
def __init__(self, val):
self.val = val
def get(self):
return self.val
with pytest.raises(
ValueError,
match=("InputNode handles user dynamic input the the DAG, and "
"cannot be used as args, kwargs, or other_args_to_resolve "
"in ClassNode constructor because it is not available at "
"class construction or binding time."),
):
with InputNode() as dag_input:
Actor.bind(dag_input)
def test_multi_input_func_dag(shared_ray_instance):
@ray.remote
def a(user_input):
return user_input * 2
@ray.remote
def b(user_input):
return user_input + 1
@ray.remote
def c(x, y):
return x + y
with InputNode() as dag_input:
a_ref = a.bind(dag_input)
b_ref = b.bind(dag_input)
dag = c.bind(a_ref, b_ref)
# (2*2) + (2*1)
assert ray.get(dag.execute(2)) == 7
# (3*2) + (3*1)
assert ray.get(dag.execute(3)) == 10
def test_input_attr_partial_access(shared_ray_instance):
@ray.remote
class Model:
def __init__(self, weight: int):
self.weight = weight
def forward(self, input: int):
return self.weight * input
@ray.remote
def combine(a, b, c, d=None):
if not d:
return a + b + c
else:
return a + b + c + d["deep"]["nested"]
# 1) Test default wrapping of args and kwargs into internal python object
with InputNode() as dag_input:
m1 = Model.bind(1)
m2 = Model.bind(2)
m1_output = m1.forward.bind(dag_input[0])
m2_output = m2.forward.bind(dag_input[1])
dag = combine.bind(m1_output, m2_output, dag_input.m3, dag_input.m4)
# 1*1 + 2*2 + 3 + 4 = 12
assert ray.get(dag.execute(1, 2, m3=3, m4={"deep": {"nested": 4}})) == 12
# 2) Test user passed data object as only input to the dag.execute()
class UserDataObj:
user_object_field_0: Any
user_object_field_1: Any
field_3: Any
def __init__(self, user_object_field_0: Any, user_object_field_1: Any,
field_3: Any) -> None:
self.user_object_field_0 = user_object_field_0
self.user_object_field_1 = user_object_field_1
self.field_3 = field_3
with InputNode() as dag_input:
m1 = Model.bind(1)
m2 = Model.bind(2)
m1_output = m1.forward.bind(dag_input.user_object_field_0)
m2_output = m2.forward.bind(dag_input.user_object_field_1)
dag = combine.bind(m1_output, m2_output, dag_input.field_3)
# 1*1 + 2*2 + 3
assert ray.get(dag.execute(UserDataObj(1, 2, 3))) == 8
# 3) Test user passed only one list object with regular list index accessor
with InputNode() as dag_input:
m1 = Model.bind(1)
m2 = Model.bind(2)
m1_output = m1.forward.bind(dag_input[0])
m2_output = m2.forward.bind(dag_input[1])
dag = combine.bind(m1_output, m2_output, dag_input[2])
# 1*1 + 2*2 + 3 + 4 = 12
assert ray.get(dag.execute([1, 2, 3])) == 8
# 4) Test user passed only one dict object with key str accessor
with InputNode() as dag_input:
m1 = Model.bind(1)
m2 = Model.bind(2)
m1_output = m1.forward.bind(dag_input["m1"])
m2_output = m2.forward.bind(dag_input["m2"])
dag = combine.bind(m1_output, m2_output, dag_input["m3"])
# 1*1 + 2*2 + 3 + 4 = 12
assert ray.get(dag.execute({"m1": 1, "m2": 2, "m3": 3})) == 8
with pytest.raises(
AssertionError,
match="Please only use int index or str as first-level key",
):
with InputNode() as dag_input:
m1 = Model.bind(1)
dag = m1.forward.bind(dag_input[(1, 2)])
def test_autoscaling_with_ensemble_nodes(serve_instance):
signal = SignalActor.remote()
autoscaling_config = {
"metrics_interval_s": 0.1,
"min_replicas": 0,
"max_replicas": 2,
"look_back_period_s": 0.4,
"downscale_delay_s": 30,
"upscale_delay_s": 0,
}
@serve.deployment(
_autoscaling_config=autoscaling_config,
_graceful_shutdown_timeout_s=1,
)
class Model:
def __init__(self, weight):
self.weight = weight
def forward(self, input):
return input + self.weight
@serve.deployment(
_autoscaling_config=autoscaling_config,
_graceful_shutdown_timeout_s=1,
)
def combine(value_refs):
ray.get(signal.wait.remote())
return sum(ray.get(value_refs))
with InputNode() as user_input:
model1 = Model.bind(0)
model2 = Model.bind(1)
output1 = model1.forward.bind(user_input)
output2 = model2.forward.bind(user_input)
output = combine.bind([output1, output2])
serve_dag = DAGDriver.options(
route_prefix="/my-dag",
_autoscaling_config=autoscaling_config,
_graceful_shutdown_timeout_s=1,
).bind(output)
dag_handle = serve.run(serve_dag)
controller = serve_instance._controller
assert get_num_running_replicas(controller, "Model") == 0
assert get_num_running_replicas(controller, "Model_1") == 0
assert get_num_running_replicas(controller, "combine") == 0
# upscaling
[dag_handle.predict.remote(0) for _ in range(10)]
wait_for_condition(
lambda: get_num_running_replicas(controller, DAGDriver.name) >= 1)
wait_for_condition(
lambda: get_num_running_replicas(controller, "Model") >= 1, timeout=40)
wait_for_condition(
lambda: get_num_running_replicas(controller, "Model_1") >= 1,
timeout=40)
wait_for_condition(
lambda: get_num_running_replicas(controller, "combine") >= 2,
timeout=40)
signal.send.remote()
# downscaling
wait_for_condition(
lambda: get_num_running_replicas(controller, DAGDriver.name) == 0,
timeout=60,
)
wait_for_condition(
lambda: get_num_running_replicas(controller, "Model") == 0,
timeout=60,
)
wait_for_condition(
lambda: get_num_running_replicas(controller, "Model_1") == 0,
timeout=60,
)
wait_for_condition(
lambda: get_num_running_replicas(controller, "combine") == 0,
timeout=60)
def test_autoscaling_with_chain_nodes(min_replicas, serve_instance):
signal = SignalActor.remote()
autoscaling_config = {
"metrics_interval_s": 0.1,
"min_replicas": min_replicas,
"max_replicas": 2,
"look_back_period_s": 0.4,
"downscale_delay_s": 30,
"upscale_delay_s": 0,
}
@serve.deployment(
autoscaling_config=autoscaling_config,
_graceful_shutdown_timeout_s=1,
)
class Model1:
def __init__(self, weight):
self.weight = weight
def forward(self, input):
ray.get(signal.wait.remote())
return input + self.weight
@serve.deployment(
autoscaling_config=autoscaling_config,
_graceful_shutdown_timeout_s=1,
)
class Model2:
def __init__(self, weight):
self.weight = weight
def forward(self, input):
return input + self.weight
with InputNode() as user_input:
model1 = Model1.bind(0)
model2 = Model2.bind(1)
output = model1.forward.bind(user_input)
output2 = model2.forward.bind(output)
serve_dag = DAGDriver.options(
route_prefix="/my-dag",
autoscaling_config=autoscaling_config,
_graceful_shutdown_timeout_s=1,
).bind(output2)
dag_handle = serve.run(serve_dag)
controller = serve_instance._controller
# upscaling
[dag_handle.predict.remote(0) for _ in range(10)]
wait_for_condition(
lambda: get_num_running_replicas(controller, DAGDriver.name) >= 1)
[dag_handle.predict.remote(0) for _ in range(10)]
wait_for_condition(
lambda: get_num_running_replicas(controller, DAGDriver.name) >= 2)
wait_for_condition(
lambda: get_num_running_replicas(controller, Model1.name) >= 1,
timeout=40)
wait_for_condition(
lambda: get_num_running_replicas(controller, Model1.name) >= 2,
timeout=40)
signal.send.remote()
wait_for_condition(
lambda: get_num_running_replicas(controller, Model2.name) >= 1,
timeout=40)
# downscaling
wait_for_condition(
lambda: get_num_running_replicas(controller, DAGDriver.name) ==
min_replicas,
timeout=60,
)
wait_for_condition(
lambda: get_num_running_replicas(controller, Model1.name) ==
min_replicas,
timeout=60,
)
wait_for_condition(
lambda: get_num_running_replicas(controller, Model2.name) ==
min_replicas,
timeout=60,
)
def get_simple_func_dag():
with InputNode() as dag_input:
ray_dag = combine.bind(dag_input[0], dag_input[1], kwargs_output=1)
return ray_dag, dag_input
ray.init()
@serve.deployment
class Model:
def __init__(self, weight):
self.weight = weight
def forward(self, input):
return input + self.weight
@serve.deployment
def combine(value1, value2, combine_type):
if combine_type == "sum":
return sum([value1, value2])
else:
return max([value1, value2])
with InputNode() as user_input:
model1 = Model.bind(0)
model2 = Model.bind(1)
output1 = model1.forward.bind(user_input[0])
output2 = model2.forward.bind(user_input[0])
dag = combine.bind(output1, output2, user_input[1])
print(ray.get(dag.execute(1, "max")))
print(ray.get(dag.execute(1, "sum")))
def get_simple_class_with_class_method_dag():
with InputNode() as dag_input:
model = Model.bind(2, ratio=0.3)
ray_dag = model.forward.bind(dag_input)
return ray_dag, dag_input