Repository: ray-project/ray
Issue category: Bug Report

Issue:
[adag] Avoid deserialization during CompiledDAGRef's deallocation
### What happened + What you expected to happen

Although we don't call ray.get, ray.get is called and deserialization still happens when the dag ref is deallocated because of the following code. https://github.com/ray-project/ray/blob/4aea49f5301ac083aed5d016919482afd2f54f83/python/ray/experimental/compiled_dag_ref.py#L80


### Versions / Dependencies

ray master

### Reproduction script

```
a = Actor.remote(0)
with InputNode() as inp:
    dag = a.inc.bind(inp)

compiled_dag = dag.experimental_compile()
ref = compiled_dag.execute(1)
compiled_dag.teardown()

# ref.get() is called upon ref's deallocation -- deserialization still happens
```

### Issue Severity

None


Hints:
@jeffreyjeffreywang would you be interested in taking this? 
@rkooo567 Yup, definitely!

Hey @stephanie-wang @ruisearch42, I just wanted to clarify a few things before I proceed with this issue. You've suggested to release the value of a CompiledDAGRef if `ray.get()` isn't called by the user (https://github.com/ray-project/ray/pull/45951#issuecomment-2168530073). **Could you please help me understand how the following code avoids execution result leak?** https://github.com/ray-project/ray/blob/44dd9a7abd0a153dd11167c8d48cde9b0a9ecfb9/python/ray/experimental/compiled_dag_ref.py#L87-L89

With this code, another issue arises -- attempting to load a python library during program exit (when CompiledDAGRef is destructed) will fail. Please refer to https://github.com/ray-project/ray/pull/47305#discussion_r1738029669 for more context.
```
foo = Foo.remote()
bar = Bar.remote()

with InputNode() as inp:
    dag = MultiOutputNode([foo.increment.bind(inp), bar.decrement.bind(inp)])

dag = dag.experimental_compile()

ref1 = dag.execute(1)
ref2 = dag.execute(1)

assert ref1.get() == [1, -1]

dag.teardown()
# Upon destruction, the DAG will be executed until the latest index.
# However, we attempt to import DAGContext during program exit. 
```

I'm thinking about removing the custom destructor entirely but wanted to understand the implications before doing so.

> With this code, another issue arises -- attempting to load a python library during program exit (when CompiledDAGRef is destructed) will fail. Please refer to https://github.com/ray-project/ray/pull/47305#discussion_r1738029669 for more context.

I think this is because python cannot guarantee all modules exist when `__del__` is called. So I agreed doing clean up this way in a destructor is a bad idea, and it is something we should clean up. 

> Could you please help me understand how the following code avoids execution result leak?

With this code, the python object is retrieved and then immediately goes out of scope. If there are any native buffers underneath, they will also be released.
> With this code, the python object is retrieved and then immediately goes out of scope. If there are any native buffers underneath, they will also be released.

Thank you @ruisearch42, could you give me an example/repro when native buffers are used and therefore this destruction is necessary? I'd like to measure whether the deserialization is necessary.

If deserialization is necessary, we still need to solve the module import issue. We might want to move the deserialization (stepping through the remaining steps in the DAG) to `teardown()`.

This is a duplicate of #46909. Will close both bugs once this issue is addressed. 
Thanks, @jeffreyjeffreywang for the great questions! The deserialization is necessary because the native buffer is reused for future data. If the reader does not explicitly read and release the buffer, then the buffer cannot be reused for future values. You can reproduce it by returning a numpy array as the DAG output; since numpy arrays are zero-copy, the buffer will be held until the np array in python goes out of scope.

Note that you do not need to deserialize the data in order to release the buffer. We just need to make sure to call the [`ReadAcquire`](https://github.com/ray-project/ray/blob/35b453445b49ad286c08099953c000c3942acf6a/src/ray/core_worker/experimental_mutable_object_manager.h#L180) and `ReadRelease` methods on the buffer (but skip the python-based deserialization).

We do a similar custom destructor for when ObjectRefs and actors go out of scope, so I think you can reuse a similar codepath to avoid the destruction ordering problem, see [here](https://github.com/ray-project/ray/blob/35b453445b49ad286c08099953c000c3942acf6a/python/ray/actor.py#L1380).
Thank you, Stephanie, for the thorough explanation. I'll dig a bit deeper and publish a PR.
Thanks @jeffreyjeffreywang ! Btw, are you in OSS ray slack? We have regular sync up, and you are more than welcome to join! 
@rkooo567 Yeah, I just joined couple days ago. Thanks for inviting, I'll keep an eye on the sync up next time and hopefully I'll be able to join! 😄 
@anyscalesam can you make sure @jeffreyjeffreywang is invited to next sync?! Thank you!

Patch:
diff --git a/python/ray/_private/worker.py b/python/ray/_private/worker.py
index 3a895356fed8..f7bbd1555ba2 100644
--- a/python/ray/_private/worker.py
+++ b/python/ray/_private/worker.py
@@ -853,6 +853,7 @@ def get_objects(
         object_refs: list,
         timeout: Optional[float] = None,
         return_exceptions: bool = False,
+        skip_deserialization: bool = False,
     ):
         """Get the values in the object store associated with the IDs.
 
@@ -869,8 +870,10 @@ def get_objects(
                 Exception object, whether to return them as values in the
                 returned list. If False, then the first found exception will be
                 raised.
+            skip_deserialization: If true, only the buffer will be released and
+                the object associated with the buffer will not be deserailized.
         Returns:
-            list: List of deserialized objects
+            list: List of deserialized objects or None if skip_deserialization is True.
             bytes: UUID of the debugger breakpoint we should drop
                 into or b"" if there is no breakpoint.
         """
@@ -903,6 +906,9 @@ def get_objects(
                     debugger_breakpoint = metadata_fields[1][
                         len(ray_constants.OBJECT_METADATA_DEBUG_PREFIX) :
                     ]
+        if skip_deserialization:
+            return None, debugger_breakpoint
+
         values = self.deserialize_objects(data_metadata_pairs, object_refs)
         if not return_exceptions:
             # Raise exceptions instead of returning them to the user.
diff --git a/python/ray/dag/compiled_dag_node.py b/python/ray/dag/compiled_dag_node.py
index 522e0779f2eb..1a4d7a7abe08 100644
--- a/python/ray/dag/compiled_dag_node.py
+++ b/python/ray/dag/compiled_dag_node.py
@@ -2056,6 +2056,21 @@ def _get_execution_results(
                 del self._result_buffer[execution_index]
         return result
 
+    def release_output_channel_buffers(self, execution_index: int):
+        from ray.dag import DAGContext
+
+        ctx = DAGContext.get_current()
+        timeout = ctx.retrieval_timeout
+
+        while self._max_finished_execution_index < execution_index:
+            self.increment_max_finished_execution_index()
+            start_time = time.monotonic()
+            self._dag_output_fetcher.release_channel_buffers(timeout)
+
+            if timeout != -1:
+                timeout -= time.monotonic() - start_time
+                timeout = max(timeout, 0)
+
     def _execute_until(
         self,
         execution_index: int,
diff --git a/python/ray/experimental/channel/common.py b/python/ray/experimental/channel/common.py
index d5f42178e1c9..9545c70c1883 100644
--- a/python/ray/experimental/channel/common.py
+++ b/python/ray/experimental/channel/common.py
@@ -342,6 +342,14 @@ def _read_list(self, timeout: Optional[float] = None) -> List[Any]:
                 timeout = max(timeout, 0)
         return results
 
+    def release_channel_buffers(self, timeout: Optional[float] = None) -> None:
+        for c in self._input_channels:
+            start_time = time.monotonic()
+            c.release_buffer(timeout)
+            if timeout is not None:
+                timeout -= time.monotonic() - start_time
+                timeout = max(timeout, 0)
+
 
 @DeveloperAPI
 class AwaitableBackgroundReader(ReaderInterface):
diff --git a/python/ray/experimental/channel/shared_memory_channel.py b/python/ray/experimental/channel/shared_memory_channel.py
index d9ada000df10..773f85a69f75 100644
--- a/python/ray/experimental/channel/shared_memory_channel.py
+++ b/python/ray/experimental/channel/shared_memory_channel.py
@@ -503,6 +503,18 @@ def read(self, timeout: Optional[float] = None) -> Any:
 
         return ret
 
+    def release_buffer(self, timeout: Optional[float] = None) -> None:
+        assert (
+            timeout is None or timeout >= 0 or timeout == -1
+        ), "Timeout must be non-negative or -1."
+        self.ensure_registered_as_reader()
+        self._worker.get_objects(
+            [self._local_reader_ref],
+            timeout=timeout,
+            return_exceptions=True,
+            skip_deserialization=True,
+        )
+
     def close(self) -> None:
         """
         Close this channel by setting the error bit on both the writer_ref and the
@@ -604,6 +616,21 @@ def read(self, timeout: Optional[float] = None) -> Any:
         self._next_read_index %= self._num_shm_buffers
         return output
 
+    def release_buffer(self, timeout: Optional[float] = None):
+        """Release the native buffer of the channel to allow the buffer to be reused for
+        future data.
+
+        If the next buffer is available, it returns immediately. If the next
+        buffer is not written by an upstream producer, it blocks until a buffer is
+        available to be released. If a buffer is not available within timeout, it raises
+        RayChannelTimeoutError.
+        """
+        # A single channel is not supposed to read and write at the same time.
+        assert self._next_write_index == 0
+        self._buffers[self._next_read_index].release_buffer(timeout)
+        self._next_read_index += 1
+        self._next_read_index %= self._num_shm_buffers
+
     def close(self) -> None:
         for buffer in self._buffers:
             buffer.close()
@@ -732,16 +759,23 @@ def write(self, value: Any, timeout: Optional[float] = None) -> None:
 
     def read(self, timeout: Optional[float] = None) -> Any:
         self.ensure_registered_as_reader()
+        return self._channel_dict[self._resolve_actor_id()].read(timeout)
+
+    def release_buffer(self, timeout: Optional[float] = None):
+        self.ensure_registered_as_reader()
+        self._channel_dict[self._resolve_actor_id()].release_buffer(timeout)
+
+    def _resolve_actor_id(self) -> str:
         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, 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)
+        return actor_id
 
     def close(self) -> None:
         for channel in self._channels:
diff --git a/python/ray/experimental/compiled_dag_ref.py b/python/ray/experimental/compiled_dag_ref.py
index da53c68d8f8a..4032f8bccc2f 100644
--- a/python/ray/experimental/compiled_dag_ref.py
+++ b/python/ray/experimental/compiled_dag_ref.py
@@ -88,9 +88,11 @@ def __del__(self):
         if self._dag.is_teardown:
             return
 
-        # If not yet, get the result and discard to avoid execution result leak.
+        # If not yet, release native buffers to avoid execution result leak. Note that
+        # we skip python-based deserialization as the values stored in the buffers are
+        # not used.
         if not self._ray_get_called:
-            self.get()
+            self._dag.release_output_channel_buffers(self._execution_index)
 
     def get(self, timeout: Optional[float] = None):
         if self._ray_get_called:


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 c764442a6c44..e457d8e9b95b 100644
--- a/python/ray/dag/tests/experimental/test_accelerated_dag.py
+++ b/python/ray/dag/tests/experimental/test_accelerated_dag.py
@@ -170,6 +170,37 @@ def test_basic(ray_start_regular):
         del result
 
 
+def test_basic_destruction(ray_start_regular):
+    a = Actor.remote(0)
+    with InputNode() as i:
+        dag = a.echo.bind(i)
+
+    compiled_dag = dag.experimental_compile()
+
+    try:
+        for i in range(3):
+            val = np.ones(100) * i
+            ref = compiled_dag.execute(val)
+            # Since ref.get() is not called, the destructor releases its native
+            # buffer without deserializing the value. If the destructor fails to
+            # release the buffer, the subsequent DAG execution will fail due to
+            # memory leak.
+            del ref
+    except RayChannelTimeoutError:
+        pytest.fail(
+            "The native buffer associated with the CompiledDAGRef was not "
+            "released upon destruction."
+        )
+
+    # Ensure that subsequent DAG executions do not fail due to memory leak
+    # and the results can be retrieved by ray.get().
+    val = np.ones(100)
+    ref = compiled_dag.execute(val)
+    result = ray.get(ref)
+    assert (result == val).all()
+    del ref
+
+
 @pytest.mark.parametrize("single_fetch", [True, False])
 def test_two_returns_one_reader(ray_start_regular, single_fetch):
     a = Actor.remote(0)