[Core][aDAG] support multi readers in multi node when dag is created …

…from an actor (ray-project#47601)

Currently, when a DAG is created from an actor, we are using different mechanism from a driver. In a driver we create a ProxyActor vs actor we are just using the actor itself.

This inconsistent mechanism is prone to error. As an example, I found when we support multi reader in multi node, we have deadlock because the driver actor needs to call ray.get(a.allocate_channel.remote()) for a downstream actor while the downstream actor calls ray.get(driver_actor.create_ref.remote()).

This fixes the issue by making ProxyActor as the default mechanism even when a dag is created inside an actor.

Signed-off-by: ujjawal-khare <ujjawal.khare@dream11.com>

python/ray/dag/compiled_dag_node.py

@@ -1105,18 +1105,16 @@ def _get_or_compile(
                         if task not in output_to_readers:
                             output_to_readers[task] = []
                         output_to_readers[task].append(downstream_task)
-                # `output_to_reader_and_node_list` stores the reader actors and
-                # the reader node IDs for each output of the current node.
-                output_to_reader_and_node_list: Dict[
-                    CompiledTask, List[Tuple["ray.actor.ActorHandle", str]]
-                ] = {task: [] for task in output_to_readers}
+                fn = task.dag_node._get_remote_method("__ray_call__")
                 for output, readers in output_to_readers.items():
                     dag_nodes = [reader.dag_node for reader in readers]
                     read_by_multi_output_node = False
                     for dag_node in dag_nodes:
                         if isinstance(dag_node, MultiOutputNode):
                             read_by_multi_output_node = True
                             break
+
+                    reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]] = []
                     if read_by_multi_output_node:
                         if len(readers) != 1:
                             raise ValueError(
@@ -1148,22 +1146,15 @@ def _get_or_compile(
 
                         #     compiled_dag.teardown()
 
-                        assert self._creator_or_proxy_actor is not None
-                        reader_and_node_list: List[
-                            Tuple["ray.actor.ActorHandle", str]
-                        ] = []
+                        assert self._proxy_actor is not None
                         for reader in readers:
                             reader_and_node_list.append(
                                 (
-                                    self._creator_or_proxy_actor,
-                                    self._get_node_id(self._creator_or_proxy_actor),
+                                    self._proxy_actor,
+                                    self._get_node_id(self._proxy_actor),
                                 )
                             )
-                        output_to_reader_and_node_list[output] = reader_and_node_list
                     else:
-                        reader_and_node_list: List[
-                            Tuple["ray.actor.ActorHandle", str]
-                        ] = []
                         # Use reader_handles_set to deduplicate readers on the
                         # same actor, because with CachedChannel each actor will
                         # only read from the upstream channel once.
@@ -1175,18 +1166,14 @@ def _get_or_compile(
                                     (reader_handle, self._get_node_id(reader_handle))
                                 )
                             reader_handles_set.add(reader_handle)
-                        output_to_reader_and_node_list[output] = reader_and_node_list
-                fn = task.dag_node._get_remote_method("__ray_call__")
-                for (
-                    output,
-                    reader_and_node_list,
-                ) in output_to_reader_and_node_list.items():
+
                     # Create an output channel for each output of the current node.
                     output_channel = ray.get(
                         fn.remote(
                             do_allocate_channel,
                             reader_and_node_list,
-                            typ=type_hint,
+                            type_hint,
+                            read_by_multi_output_node,
                         )
                     )
                     output_idx = None
@@ -1238,7 +1225,8 @@ def _get_or_compile(
                     do_allocate_channel(
                         self,
                         reader_and_node_list,
-                        typ=type_hint,
+                        type_hint,
+                        False,
                     )
                 ]
             else:
@@ -1288,7 +1276,6 @@ def _get_or_compile(
         assert len(input_task.output_channels) == 1
         self.dag_input_channel = input_task.output_channels[0]
         assert self.dag_input_channel is not None
-
         # Create executable tasks for each actor
         for actor_handle, tasks in self.actor_to_tasks.items():
             # Dict from non-dag-input arg to the set of tasks that consume it.

python/ray/dag/tests/experimental/test_multi_node_dag.py

@@ -298,6 +298,50 @@ def _get_node_id(self) -> "ray.NodeID":
     compiled_dag.teardown()
 
 
+def test_multi_node_dag_from_actor(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=1)
+    ray.init()
+    cluster.add_node(num_cpus=1)
+
+    @ray.remote(num_cpus=0)
+    class SameNodeActor:
+        def predict(self, x: str):
+            return x
+
+    @ray.remote(num_cpus=1)
+    class RemoteNodeActor:
+        def predict(self, x: str, y: str):
+            return y
+
+    @ray.remote(num_cpus=1)
+    class DriverActor:
+        def __init__(self):
+            self._base_actor = SameNodeActor.options(
+                scheduling_strategy=NodeAffinitySchedulingStrategy(
+                    ray.get_runtime_context().get_node_id(), soft=False
+                )
+            ).remote()
+            self._refiner_actor = RemoteNodeActor.remote()
+
+            with InputNode() as inp:
+                x = self._base_actor.predict.bind(inp)
+                dag = self._refiner_actor.predict.bind(
+                    inp,
+                    x,
+                )
+
+            self._adag = dag.experimental_compile(
+                _execution_timeout=120,
+            )
+
+        def call(self, prompt: str) -> bytes:
+            return ray.get(self._adag.execute(prompt))
+
+    parallel = DriverActor.remote()
+    assert ray.get(parallel.call.remote("abc")) == "abc"
+
+
 if __name__ == "__main__":
     if os.environ.get("PARALLEL_CI"):
         sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))

python/ray/experimental/channel/shared_memory_channel.py

@@ -157,7 +157,12 @@ def create_channel(
                     cpu_data_typ=cpu_data_typ,
                 )
 
-        return CompositeChannel(writer, reader_and_node_list, self._num_shm_buffers)
+        return CompositeChannel(
+            writer,
+            reader_and_node_list,
+            self._num_shm_buffers,
+            read_by_adag_driver,
+        )
 
     def set_nccl_group_id(self, group_id: str) -> None:
         assert self.requires_nccl()
@@ -652,6 +657,7 @@ def __init__(
         writer: Optional[ray.actor.ActorHandle],
         reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
         num_shm_buffers: int,
+        read_by_adag_driver: bool,
         _channel_dict: Optional[Dict[ray.ActorID, ChannelInterface]] = None,
         _channels: Optional[Set[ChannelInterface]] = None,
         _writer_registered: bool = False,
@@ -738,6 +744,7 @@ def __reduce__(self):
             self._writer,
             self._reader_and_node_list,
             self._num_shm_buffers,
+            self._read_by_adag_driver,
             self._channel_dict,
             self._channels,
             self._writer_registered,

python/ray/tests/test_channel.py

@@ -906,8 +906,12 @@ def __init__(self):
         def pass_channel(self, channel):
             self._chan = channel
 
-        def create_composite_channel(self, writer, reader_and_node_list):
-            self._chan = ray_channel.CompositeChannel(writer, reader_and_node_list, 10)
+        def create_composite_channel(
+            self, writer, reader_and_node_list, read_by_adag_driver
+        ):
+            self._chan = ray_channel.CompositeChannel(
+                writer, reader_and_node_list, 10, read_by_adag_driver
+            )
             return self._chan
 
         def read(self):
@@ -922,7 +926,9 @@ def write(self, value):
     node2 = get_actor_node_id(actor2)
 
     # Create a channel to communicate between driver process and actor1.
-    driver_to_actor1_channel = ray_channel.CompositeChannel(None, [(actor1, node1)], 10)
+    driver_to_actor1_channel = ray_channel.CompositeChannel(
+        None, [(actor1, node1)], 10, False
+    )
     ray.get(actor1.pass_channel.remote(driver_to_actor1_channel))
     driver_to_actor1_channel.write("hello")
     assert ray.get(actor1.read.remote()) == "hello"
@@ -979,7 +985,9 @@ def pass_channel(self, channel):
             self._chan = channel
 
         def create_composite_channel(self, writer, reader_and_node_list):
-            self._chan = ray_channel.CompositeChannel(writer, reader_and_node_list, 10)
+            self._chan = ray_channel.CompositeChannel(
+                writer, reader_and_node_list, 10, False
+            )
             return self._chan
 
         def read(self):
@@ -995,7 +1003,7 @@ def write(self, value):
 
     # The driver writes data to CompositeChannel and actor1 and actor2 read it.
     driver_output_channel = ray_channel.CompositeChannel(
-        None, [(actor1, node1), (actor2, node2)], 10
+        None, [(actor1, node1), (actor2, node2)], 10, False
     )
     ray.get(actor1.pass_channel.remote(driver_output_channel))
     ray.get(actor2.pass_channel.remote(driver_output_channel))