[core][experimental] Add experimental compiled DAG

python/ray/_private/ray_perf.py

@@ -9,6 +9,7 @@
 import ray
 
 import ray.experimental.channel as ray_channel
+from ray.dag import InputNode, MultiOutputNode
 
 logger = logging.getLogger(__name__)
 
@@ -369,6 +370,58 @@ def read(self, chans):
     for reader in readers:
         ray.kill(reader)
 
+    # Tests for compiled DAGs.
+
+    def _exec(dag):
+        output_channel = dag.execute(b"x")
+        output_channel.begin_read()
+        output_channel.end_read()
+
+    def _exec_multi_output(dag):
+        output_channels = dag.execute(b"x")
+        for output_channel in output_channels:
+            output_channel.begin_read()
+        for output_channel in output_channels:
+            output_channel.end_read()
+
+    @ray.remote
+    class Actor:
+        def echo(self, x):
+            return x
+
+    a = Actor.remote()
+    with InputNode() as inp:
+        dag = a.echo.bind(inp)
+
+    results += timeit("single-actor DAG calls", lambda: ray.get(dag.execute(b"x")))
+    dag = dag.experimental_compile()
+    results += timeit("compiled single-actor DAG calls", lambda: _exec(dag))
+
+    del a
+    n_cpu = multiprocessing.cpu_count() // 2
+    actors = [Actor.remote() for _ in range(n_cpu)]
+    with InputNode() as inp:
+        dag = MultiOutputNode([a.echo.bind(inp) for a in actors])
+    results += timeit(
+        "scatter-gather DAG calls, n={n_cpu} actors", lambda: ray.get(dag.execute(b"x"))
+    )
+    dag = dag.experimental_compile()
+    results += timeit(
+        f"compiled scatter-gather DAG calls, n={n_cpu} actors",
+        lambda: _exec_multi_output(dag),
+    )
+
+    actors = [Actor.remote() for _ in range(n_cpu)]
+    with InputNode() as inp:
+        dag = inp
+        for a in actors:
+            dag = a.echo.bind(dag)
+    results += timeit(
+        f"chain DAG calls, n={n_cpu} actors", lambda: ray.get(dag.execute(b"x"))
+    )
+    dag = dag.experimental_compile()
+    results += timeit(f"compiled chain DAG calls, n={n_cpu} actors", lambda: _exec(dag))
+
     ray.shutdown()
 
     ############################

python/ray/dag/BUILD

@@ -68,3 +68,11 @@ py_test(
     tags = ["exclusive", "team:core", "ray_dag_tests"],
     deps = [":dag_lib"],
 )
+
+py_test(
+    name = "test_accelerated_dag",
+    size = "medium",
+    srcs = dag_tests_srcs,
+    tags = ["exclusive", "team:core", "ray_dag_tests"],
+    deps = [":dag_lib"],
+)

python/ray/dag/class_node.py

@@ -202,3 +202,12 @@ def __str__(self) -> str:
 
     def get_method_name(self) -> str:
         return self._method_name
+
+    def _get_remote_method(self, method_name):
+        method_body = getattr(self._parent_class_node, method_name)
+        return method_body
+
+    def _get_actor_handle(self) -> Optional["ray.actor.ActorHandle"]:
+        if not isinstance(self._parent_class_node, ray.actor.ActorHandle):
+            return None
+        return self._parent_class_node