Repository: ray-project/ray
Issue category: Feature Request

Issue:
[core][adag] Support reading the same node from another DAG node and the driver
### What happened + What you expected to happen

Currently a DAG node can only be read by other DAG nodes or by the driver. We should support allowing it to be read by both. This should already work since the driver is represented by a proxy actor, so there might just be too strong of an assertion check right now.

### Versions / Dependencies

3.0dev

### Reproduction script

```python
with InputNode() as inp:
  x = a.foo.bind(inp)
  y = b.foo.bind(x)
  dag = MultiOutputNode([x, y])

dag = dag.experimental_compile()
```

### Issue Severity

None


Hints:
This is duplicate with https://github.com/ray-project/ray/issues/47041
I don't think this is the same as #47041? This issue is reading the same channel from different actors, not reading the same channel multiple times from the same actor?
Oops, sorry. I don't read it carefully.

Patch:
diff --git a/python/ray/dag/compiled_dag_node.py b/python/ray/dag/compiled_dag_node.py
index 7c6160d8937d..b8fce79a123c 100644
--- a/python/ray/dag/compiled_dag_node.py
+++ b/python/ray/dag/compiled_dag_node.py
@@ -101,7 +101,7 @@ def do_allocate_channel(
     self,
     reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
     typ: ChannelOutputType,
-    read_by_adag_driver: bool,
+    driver_actor_id: Optional[str] = None,
 ) -> ChannelInterface:
     """Generic actor method to allocate an output channel.
 
@@ -109,8 +109,8 @@ def do_allocate_channel(
         reader_and_node_list: A list of tuples, where each tuple contains a reader
             actor handle and the node ID where the actor is located.
         typ: The output type hint for the channel.
-        read_by_adag_driver: True if the channel will be read by an aDAG driver
-            (Ray driver or actor and task that creates an aDAG).
+        driver_actor_id: If this channel is read by a driver and that driver is an
+            actual actor, this will be the actor ID of that driver actor.
 
     Returns:
         The allocated channel.
@@ -126,7 +126,7 @@ def do_allocate_channel(
     output_channel = typ.create_channel(
         writer,
         reader_and_node_list,
-        read_by_adag_driver,
+        driver_actor_id,
     )
     return output_channel
 
@@ -1312,73 +1312,46 @@ def _get_or_compile(
                         output_to_readers[task].append(downstream_task)
                 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]
+                    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.
+                    reader_handles_set = set()
                     read_by_multi_output_node = False
-                    for dag_node in dag_nodes:
-                        if isinstance(dag_node, MultiOutputNode):
+                    for reader in readers:
+                        if isinstance(reader.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(
-                                "DAG outputs currently can only be read by the "
-                                "driver or the same actor that is also the "
-                                "InputNode, not by both the driver and actors.",
-                            )
-
-                        # This node is a multi-output node, which means it will
-                        # only be read by the driver or the actor that is also
-                        # the InputNode.
-
-                        # TODO(jhumphri): Handle case where there is an actor,
-                        # other than just the driver actor, also reading the
-                        # output from the `task` node. For example, the following
-                        # currently does not work:
-                        # def test_blah(ray_start_regular):
-                        #     a = Actor.remote(0)
-                        #     b = Actor.remote(10)
-                        #     with InputNode() as inp:
-                        #         x = a.inc.bind(inp)
-                        #         y = b.inc.bind(x)
-                        #         dag = MultiOutputNode([x, y])
-
-                        #     compiled_dag = dag.experimental_compile()
-                        #     output_channel = compiled_dag.execute(1)
-                        #     result = output_channel.read()
-                        #     print(result)
-
-                        #     compiled_dag.teardown()
-
-                        assert self._proxy_actor is not None
-                        for reader in readers:
-                            reader_and_node_list.append(
+                            # inserting at 0 to make sure driver is first reader as
+                            # expected by CompositeChannel read
+                            reader_and_node_list.insert(
+                                0,
                                 (
                                     self._proxy_actor,
                                     self._get_node_id(self._proxy_actor),
-                                )
+                                ),
                             )
-                    else:
-                        # Use reader_handles_set to deduplicate readers on the
-                        # same actor, because with CachedChannel each actor will
-                        # only read from the upstream channel once.
-                        reader_handles_set = set()
-                        for reader in readers:
+                        else:
                             reader_handle = reader.dag_node._get_actor_handle()
                             if reader_handle not in reader_handles_set:
+                                reader_handle = reader.dag_node._get_actor_handle()
                                 reader_and_node_list.append(
                                     (reader_handle, self._get_node_id(reader_handle))
                                 )
-                            reader_handles_set.add(reader_handle)
+                                reader_handles_set.add(reader_handle)
 
+                    # if driver is an actual actor, gets driver actor id
+                    driver_actor_id = (
+                        ray.get_runtime_context().get_actor_id()
+                        if read_by_multi_output_node
+                        else None
+                    )
                     # Create an output channel for each output of the current node.
                     output_channel = ray.get(
                         fn.remote(
                             do_allocate_channel,
                             reader_and_node_list,
                             type_hint,
-                            read_by_multi_output_node,
+                            driver_actor_id,
                         )
                     )
                     output_idx = None
@@ -1444,7 +1417,7 @@ def _get_or_compile(
                         self,
                         reader_and_node_list,
                         type_hint,
-                        False,
+                        None,
                     )
                     task.output_channels.append(output_channel)
                     task.output_idxs.append(
diff --git a/python/ray/experimental/channel/common.py b/python/ray/experimental/channel/common.py
index 55715eefc856..44ccd2b23b27 100644
--- a/python/ray/experimental/channel/common.py
+++ b/python/ray/experimental/channel/common.py
@@ -82,7 +82,7 @@ def create_channel(
         self,
         writer: Optional["ray.actor.ActorHandle"],
         reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
-        read_by_adag_driver: bool,
+        driver_actor_id: Optional[str] = None,
     ) -> "ChannelInterface":
         """
         Instantiate a ChannelInterface class that can be used
@@ -92,8 +92,9 @@ def create_channel(
             writer: The actor that may write to the channel. None signifies the driver.
             reader_and_node_list: A list of tuples, where each tuple contains a reader
                 actor handle and the node ID where the actor is located.
-            read_by_adag_driver: True if a channel is read by an aDAG driver (Ray driver
-                or an actor that creates the aDAG).
+            driver_actor_id: If this is a CompositeChannel that is read by a driver and
+                that driver is an actual actor, this will be the actor ID of that
+                driver actor.
         Returns:
             A ChannelInterface that can be used to pass data
                 of this type.
diff --git a/python/ray/experimental/channel/shared_memory_channel.py b/python/ray/experimental/channel/shared_memory_channel.py
index fd98d95b2d4a..d9ada000df10 100644
--- a/python/ray/experimental/channel/shared_memory_channel.py
+++ b/python/ray/experimental/channel/shared_memory_channel.py
@@ -130,7 +130,7 @@ def create_channel(
         self,
         writer: Optional["ray.actor.ActorHandle"],
         reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
-        read_by_adag_driver: bool,
+        driver_actor_id: Optional[str] = None,
     ) -> "Channel":
         """
         Instantiate a ChannelInterface class that can be used
@@ -140,8 +140,8 @@ def create_channel(
             writer: The actor that may write to the channel. None signifies the driver.
             reader_and_node_list: A list of tuples, where each tuple contains a reader
                 actor handle and the node ID where the actor is located.
-        read_by_adag_driver: True if the channel will be read by an aDAG driver
-            (Ray driver or actor and task that creates an aDAG).
+            driver_actor_id: If this channel is read by a driver and that driver is an
+                actual actor, this will be the actor ID of that driver actor.
 
         Returns:
             A ChannelInterface that can be used to pass data
@@ -151,7 +151,7 @@ def create_channel(
             writer,
             reader_and_node_list,
             self._num_shm_buffers,
-            read_by_adag_driver,
+            driver_actor_id,
         )
 
 
@@ -631,8 +631,8 @@ class CompositeChannel(ChannelInterface):
         writer: The actor that may write to the channel. None signifies the driver.
         reader_and_node_list: A list of tuples, where each tuple contains a reader
             actor handle and the node ID where the actor is located.
-        read_by_adag_driver: True if the channel will be read by a driver (Ray driver or
-            actor and task that creates an aDAG.
+        driver_actor_id: If this channel is read by a driver and that driver is an
+            actual actor, this will be the actor ID of that driver actor.
     """
 
     def __init__(
@@ -640,7 +640,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,
+        driver_actor_id: Optional[str] = None,
         _channel_dict: Optional[Dict[ray.ActorID, ChannelInterface]] = None,
         _channels: Optional[Set[ChannelInterface]] = None,
         _writer_registered: bool = False,
@@ -649,13 +649,13 @@ def __init__(
         self._writer = writer
         self._reader_and_node_list = reader_and_node_list
         self._num_shm_buffers = num_shm_buffers
+        self._driver_actor_id = driver_actor_id
         self._writer_registered = _writer_registered
         self._reader_registered = _reader_registered
         # A dictionary that maps the actor ID to the channel object.
         self._channel_dict = _channel_dict or {}
         # The set of channels is a deduplicated version of the _channel_dict values.
         self._channels = _channels or set()
-        self._read_by_adag_driver = read_by_adag_driver
         if self._channels:
             # This CompositeChannel object is created by deserialization.
             # We don't need to create channels again.
@@ -693,21 +693,6 @@ def __init__(
     def _get_actor_id(self, reader: ray.actor.ActorHandle) -> str:
         return reader._actor_id.hex()
 
-    def _get_self_actor_id(self) -> str:
-        """
-        Get the actor ID of the current process. If the current process is the
-        aDAG owner (e.g., driver or an actor/task that creates aDAG),
-        use the actor ID of the DAGDriverProxyActor.
-        """
-        actor_id = ray.get_runtime_context().get_actor_id()
-        if self._read_by_adag_driver:
-            # The reader is the driver process.
-            # Use the actor ID of the DAGDriverProxyActor.
-            assert len(self._reader_and_node_list) == 1
-            driver_actor = self._reader_and_node_list[0][0]
-            actor_id = self._get_actor_id(driver_actor)
-        return actor_id
-
     def ensure_registered_as_writer(self) -> None:
         if self._writer_registered:
             return
@@ -727,7 +712,7 @@ def __reduce__(self):
             self._writer,
             self._reader_and_node_list,
             self._num_shm_buffers,
-            self._read_by_adag_driver,
+            self._driver_actor_id,
             self._channel_dict,
             self._channels,
             self._writer_registered,
@@ -747,7 +732,15 @@ def write(self, value: Any, timeout: Optional[float] = None) -> None:
 
     def read(self, timeout: Optional[float] = None) -> Any:
         self.ensure_registered_as_reader()
-        actor_id = self._get_self_actor_id()
+        actor_id = ray.get_runtime_context().get_actor_id()
+        # if actor_id is None, read was called by the driver
+        # if the driver is an actor, driver_actor_id will be set to that actor id
+        if actor_id is None or actor_id == self._driver_actor_id:
+            # Use the actor ID of the DAGDriverProxyActor.
+            # The proxy actor is always the first actor in the reader_and_node_list.
+            assert len(self._reader_and_node_list) >= 1
+            driver_proxy_actor = self._reader_and_node_list[0][0]
+            actor_id = self._get_actor_id(driver_proxy_actor)
         return self._channel_dict[actor_id].read(timeout)
 
     def close(self) -> None:
diff --git a/python/ray/experimental/channel/torch_tensor_nccl_channel.py b/python/ray/experimental/channel/torch_tensor_nccl_channel.py
index 4b8baedeebd0..8a0a043a38f3 100644
--- a/python/ray/experimental/channel/torch_tensor_nccl_channel.py
+++ b/python/ray/experimental/channel/torch_tensor_nccl_channel.py
@@ -354,7 +354,7 @@ def __init__(
             self._meta_channel = metadata_type.create_channel(
                 self._writer,
                 self._reader_and_node_list,
-                False,
+                None,
             )
 
     def ensure_registered_as_writer(self):
diff --git a/python/ray/experimental/channel/torch_tensor_type.py b/python/ray/experimental/channel/torch_tensor_type.py
index 8615f18b7d65..d048a37abbb5 100644
--- a/python/ray/experimental/channel/torch_tensor_type.py
+++ b/python/ray/experimental/channel/torch_tensor_type.py
@@ -122,7 +122,7 @@ def create_channel(
         self,
         writer: Optional["ray.actor.ActorHandle"],
         reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
-        read_by_adag_driver: bool,
+        driver_actor_id: Optional[str] = None,
         _cpu_data_channel: Optional["Channel"] = None,
         _tensor_metadata_channel: Optional["Channel"] = None,
     ) -> type:
@@ -142,9 +142,7 @@ def create_channel(
             if _cpu_data_channel is None and not self._direct_return:
                 # Create a CPU channel to send non-tensor data.
                 _cpu_data_channel = SharedMemoryType().create_channel(
-                    writer,
-                    reader_and_node_list,
-                    read_by_adag_driver,
+                    writer, reader_and_node_list, driver_actor_id
                 )
 
             return TorchTensorNcclChannel(
@@ -159,7 +157,7 @@ def create_channel(
         # shared-memory channel.
         # TODO(swang): Allow the initial max buffer size to be overridden.
         typ = SharedMemoryType()
-        return typ.create_channel(writer, reader_and_node_list, read_by_adag_driver)
+        return typ.create_channel(writer, reader_and_node_list, driver_actor_id)
 
     def requires_nccl(self) -> bool:
         return self.transport == self.NCCL


Test patch:
diff --git a/python/ray/dag/tests/experimental/test_accelerated_dag.py b/python/ray/dag/tests/experimental/test_accelerated_dag.py
index cbec80a871c4..2142de8312b1 100644
--- a/python/ray/dag/tests/experimental/test_accelerated_dag.py
+++ b/python/ray/dag/tests/experimental/test_accelerated_dag.py
@@ -2078,6 +2078,18 @@ def foo(self, arg):
         dag.execute(1)
 
 
+def test_multiple_reads_from_same_actor(ray_start_cluster):
+    a = Actor.remote(0)
+    b = Actor.remote(10)
+    with InputNode() as inp:
+        x = a.inc.bind(inp)
+        y = b.inc.bind(x)
+        z = b.inc.bind(x)
+        dag = MultiOutputNode([y, z])
+    dag = dag.experimental_compile()
+    assert ray.get(dag.execute(1)) == [11, 12]
+
+
 def test_driver_and_actor_as_readers(ray_start_cluster):
     a = Actor.remote(0)
     b = Actor.remote(10)
@@ -2085,14 +2097,24 @@ def test_driver_and_actor_as_readers(ray_start_cluster):
         x = a.inc.bind(inp)
         y = b.inc.bind(x)
         dag = MultiOutputNode([x, y])
+    dag = dag.experimental_compile()
+    assert ray.get(dag.execute(1)) == [1, 11]
 
-    with pytest.raises(
-        ValueError,
-        match="DAG outputs currently can only be read by the driver or "
-        "the same actor that is also the InputNode, not by both "
-        "the driver and actors.",
-    ):
-        dag.experimental_compile()
+
+def test_driver_and_intraprocess_read(ray_start_cluster):
+    """
+    This test is similar to the `test_driver_and_actor_as_readers` test, but now for x,
+    there is IntraProcessChannel to Actor a and a BufferedSharedMemoryChannel to the
+    driver and the CompositeChannel has to choose the correct channel to read from in
+    both situations.
+    """
+    a = Actor.remote(0)
+    with InputNode() as inp:
+        x = a.inc.bind(inp)
+        y = a.inc.bind(x)
+        dag = MultiOutputNode([x, y])
+    dag = dag.experimental_compile()
+    assert ray.get(dag.execute(1)) == [1, 2]
 
 
 @pytest.mark.skip("Currently buffer size is set to 1 because of regression.")
@@ -2297,6 +2319,45 @@ def call(self, value):
     assert ray.get(ref) == 3
 
 
+def test_driver_as_actor_and_actor_reading(ray_start_cluster):
+    @ray.remote
+    class Replica:
+        def __init__(self):
+            self.w = TestWorker.remote()
+            self.w2 = TestWorker.remote()
+            with InputNode() as inp:
+                x = self.w.add_one.bind(inp)
+                y = self.w2.add_one.bind(x)
+                dag = MultiOutputNode([x, y])
+            self.compiled_dag = dag.experimental_compile()
+
+        def exec_and_get(self, value):
+            return ray.get(self.compiled_dag.execute(value))
+
+    replica = Replica.remote()
+    result = replica.exec_and_get.remote(1)
+    assert ray.get(result) == [2, 3]
+
+
+def test_driver_as_actor_and_intraprocess_read(ray_start_cluster):
+    @ray.remote
+    class Replica:
+        def __init__(self):
+            self.w = TestWorker.remote()
+            with InputNode() as inp:
+                x = self.w.add_one.bind(inp)
+                y = self.w.add_one.bind(x)
+                dag = MultiOutputNode([x, y])
+            self.compiled_dag = dag.experimental_compile()
+
+        def exec_and_get(self, value):
+            return ray.get(self.compiled_dag.execute(value))
+
+    replica = Replica.remote()
+    result = replica.exec_and_get.remote(1)
+    assert ray.get(result) == [2, 3]
+
+
 @pytest.mark.parametrize("single_fetch", [True, False])
 def test_multiple_readers_multiple_writers(shutdown_only, single_fetch):
     """
diff --git a/python/ray/tests/test_channel.py b/python/ray/tests/test_channel.py
index 55f7a6d49ee8..69127ee157a8 100644
--- a/python/ray/tests/test_channel.py
+++ b/python/ray/tests/test_channel.py
@@ -907,10 +907,10 @@ def pass_channel(self, channel):
             self._chan = channel
 
         def create_composite_channel(
-            self, writer, reader_and_node_list, read_by_adag_driver
+            self, writer, reader_and_node_list, driver_actor_id
         ):
             self._chan = ray_channel.CompositeChannel(
-                writer, reader_and_node_list, 10, read_by_adag_driver
+                writer, reader_and_node_list, 10, driver_actor_id
             )
             return self._chan
 
@@ -926,21 +926,19 @@ 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, False
-    )
+    driver_to_actor1_channel = ray_channel.CompositeChannel(None, [(actor1, node1)], 10)
     ray.get(actor1.pass_channel.remote(driver_to_actor1_channel))
     driver_to_actor1_channel.write("hello")
     assert ray.get(actor1.read.remote()) == "hello"
 
     # Create a channel to communicate between two tasks in actor1.
-    ray.get(actor1.create_composite_channel.remote(actor1, [(actor1, node1)], False))
+    ray.get(actor1.create_composite_channel.remote(actor1, [(actor1, node1)], None))
     ray.get(actor1.write.remote("world"))
     assert ray.get(actor1.read.remote()) == "world"
 
     # Create a channel to communicate between actor1 and actor2.
     actor1_to_actor2_channel = ray.get(
-        actor1.create_composite_channel.remote(actor1, [(actor2, node2)], False)
+        actor1.create_composite_channel.remote(actor1, [(actor2, node2)], None)
     )
     ray.get(actor2.pass_channel.remote(actor1_to_actor2_channel))
     ray.get(actor1.write.remote("hello world"))
@@ -950,7 +948,9 @@ def write(self, value):
     driver_actor = create_driver_actor()
     actor2_to_driver_channel = ray.get(
         actor2.create_composite_channel.remote(
-            actor2, [(driver_actor, get_actor_node_id(driver_actor))], True
+            actor2,
+            [(driver_actor, get_actor_node_id(driver_actor))],
+            driver_actor._actor_id.hex(),
         )
     )
     ray.get(actor2.write.remote("world hello"))
@@ -985,9 +985,7 @@ 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, False
-            )
+            self._chan = ray_channel.CompositeChannel(writer, reader_and_node_list, 10)
             return self._chan
 
         def read(self):
@@ -1003,7 +1001,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, False
+        None, [(actor1, node1), (actor2, node2)], 10
     )
     ray.get(actor1.pass_channel.remote(driver_output_channel))
     ray.get(actor2.pass_channel.remote(driver_output_channel))
@@ -1033,12 +1031,19 @@ def write(self, value):
         # are the same actor. Note that reading the channel multiple
         # times is supported via channel cache mechanism.
         ray.get(actor1.read.remote())
-    """
-    TODO (kevin85421): Add tests for the following cases:
-    (1) actor1 writes data to CompositeChannel and two Ray tasks on actor2 read it.
-    (2) actor1 writes data to CompositeChannel and actor2 and the driver reads it.
-    Currently, (1) is not supported, and (2) is blocked by the reference count issue.
-    """
+
+    # actor1 writes data to CompositeChannel and actor2 and the driver reads it.
+    driver_actor = create_driver_actor()
+    actor1_output_channel = ray.get(
+        actor1.create_composite_channel.remote(
+            actor1,
+            [(driver_actor, get_actor_node_id(driver_actor)), (actor2, node2)],
+        )
+    )
+    ray.get(actor2.pass_channel.remote(actor1_output_channel))
+    ray.get(actor1.write.remote("world hello"))
+    assert ray.get(actor2.read.remote()) == "world hello"
+    assert actor1_output_channel.read() == "world hello"
 
 
 @pytest.mark.skipif(
diff --git a/python/ray/tests/test_nccl_channel.py b/python/ray/tests/test_nccl_channel.py
index 919189af9ed1..d75bd1f7b637 100644
--- a/python/ray/tests/test_nccl_channel.py
+++ b/python/ray/tests/test_nccl_channel.py
@@ -69,7 +69,7 @@ def create_nccl_channel(
         self.tensor_chan = typ.create_channel(
             ray.get_runtime_context().current_actor,
             reader_and_node_list,
-            read_by_adag_driver=False,
+            driver_actor_id=None,
             _cpu_data_channel=cpu_data_channel,
             _tensor_metadata_channel=tensor_metadata_channel,
         )