chore: import upstream snapshot with attribution

python/ray/dag/tests/test_output_node.py

@@ -0,0 +1,203 @@
+import pytest
+
+import ray
+from ray._common.test_utils import wait_for_condition
+from ray.dag.input_node import InputNode
+from ray.dag.output_node import MultiOutputNode
+from ray.util.state import list_tasks
+
+
+def test_output_node(shared_ray_instance):
+    @ray.remote
+    def f(input):
+        return input
+
+    with pytest.raises(ValueError):
+        with InputNode() as input_data:
+            dag = MultiOutputNode(f.bind(input_data))
+
+    with InputNode() as input_data:
+        dag = MultiOutputNode([f.bind(input_data)])
+
+    assert ray.get(dag.execute(1)) == [1]
+    assert ray.get(dag.execute(2)) == [2]
+
+    with InputNode() as input_data:
+        dag = MultiOutputNode([f.bind(input_data["x"]), f.bind(input_data["y"])])
+
+    refs = dag.execute({"x": 1, "y": 2})
+    assert len(refs) == 2
+    assert ray.get(refs) == [1, 2]
+
+    with InputNode() as input_data:
+        dag = MultiOutputNode(
+            [f.bind(input_data["x"]), f.bind(input_data["y"]), f.bind(input_data["x"])]
+        )
+
+    refs = dag.execute({"x": 1, "y": 2})
+    assert len(refs) == 3
+    assert ray.get(refs) == [1, 2, 1]
+
+
+def test_dag_with_actor_handle(shared_ray_instance):
+    """Verify DAG API works with actor created by .remote"""
+
+    @ray.remote
+    class Worker:
+        def __init__(self):
+            self.forward_called = 0
+            self.init_called = 0
+
+        def forward(self, input):
+            print("forward")
+            self.forward_called += 1
+            return input
+
+        def initialize(self, input):
+            print("initialize")
+            self.init_called += 1
+            return input
+
+        def get(self):
+            return (self.forward_called, self.init_called)
+
+    worker = Worker.remote()
+    with InputNode() as input_node:
+        init_dag = worker.initialize.bind(input_node)
+    with InputNode() as input_node:
+        forward_dag = worker.forward.bind(input_node)
+
+    assert ray.get(init_dag.execute(1)) == 1
+    assert ray.get(forward_dag.execute(2)) == 2
+
+    # Make sure both forward/initialize called only once
+    assert ray.get(worker.get.remote()) == (1, 1)
+
+    # Double check the actor is resued.
+    assert ray.get(init_dag.execute(1)) == 1
+    assert ray.get(worker.get.remote()) == (1, 2)
+
+
+def test_dag_with_alive_actors_chained(shared_ray_instance):
+    """Verify we can have multiple DAGs to the
+    same actor that are chained.
+    """
+
+    @ray.remote
+    class Actor:
+        def __init__(self, init_value):
+            self.i = init_value
+
+        def add(self, x):
+            return self.i + x
+
+    @ray.remote
+    def combine(x, y):
+        return x + y
+
+    a1 = Actor.remote(10)
+    a1_dag = a1.add.bind(a1.add.bind(2))  # 22
+    a1_dag_2 = a1.add.bind(a1.add.bind(6))  # 26
+    dag = combine.bind(a1_dag, a1_dag_2)
+
+    assert ray.get(dag.execute()) == 48
+
+
+def test_tensor_parallel_dag(shared_ray_instance):
+    """Simulate the TP DAG with N workers.
+    Input -> forward -> MultiOutput
+    """
+
+    @ray.remote
+    class Worker:
+        def __init__(self, rank):
+            self.rank = rank
+            self.forwarded = 0
+
+        def forward(self, input_data: int):
+            print(input_data)
+            self.forwarded += 1
+            return self.rank + input_data
+
+        def initialize(self):
+            pass
+
+        def get_forwarded(self):
+            return self.forwarded
+
+    NUM_WORKERS = 4
+    workers = [Worker.remote(i) for i in range(NUM_WORKERS)]
+    # Init multiple times.
+    for _ in range(4):
+        ray.get([worker.initialize.remote() for worker in workers])
+
+    with InputNode() as input_data:
+        dag = MultiOutputNode([worker.forward.bind(input_data) for worker in workers])
+
+    # Run DAG repetitively.
+    ITER = 4
+    assert ITER > 1
+    for i in range(ITER):
+        ref = dag.execute(i)
+        all_outputs = ray.get(ref)
+        assert len(all_outputs) == NUM_WORKERS
+        assert all_outputs == [i + j for j in range(NUM_WORKERS)]
+
+    forwarded = ray.get([worker.get_forwarded.remote() for worker in workers])
+    assert forwarded == [ITER for _ in range(NUM_WORKERS)]
+
+
+def test_shared_output(shared_ray_instance):
+    """Verify when an upstream task output is shared by
+    multi output, the upstream task runs only once.
+    """
+
+    @ray.remote
+    def shared_f():
+        return 1
+
+    @ray.remote
+    def g(input):
+        return input + 1
+
+    @ray.remote
+    def h(input):
+        return input + 2
+
+    x = shared_f.bind()
+    dag = MultiOutputNode([g.bind(x), h.bind(x)])
+
+    assert ray.get(dag.execute()) == [2, 3]
+
+    # Verify f ran only once.
+    def verify():
+        tasks = list_tasks(filters=[("name", "=", "shared_f")])
+        return len(tasks) == 1
+
+    wait_for_condition(verify)
+
+
+def test_bind_survives_handle_deletion(shared_ray_instance):
+    """Verify that .bind().execute() still works even if the original handle was dropped."""
+
+    @ray.remote
+    class A:
+        def f(self):
+            return 1
+
+    # Grab the handle and the bound method node
+    actor = A.remote()
+    method_node = actor.f.bind()
+
+    # Destroy the only Python variable reference and force collection
+    del actor
+
+    # Executing should now succeed because the node holds the ref
+    result = ray.get(method_node.execute())
+    assert result == 1
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))