diff --git a/python/ray/data/_internal/numpy_support.py b/python/ray/data/_internal/numpy_support.py
index d04060fc831e8..28cb0e537fa76 100644
--- a/python/ray/data/_internal/numpy_support.py
+++ b/python/ray/data/_internal/numpy_support.py
@@ -45,28 +45,94 @@ def validate_numpy_batch(batch: Union[Dict[str, np.ndarray], Dict[str, list]]) -

 def _detect_highest_datetime_precision(datetime_list: List[datetime]) -> str:
-    highest_precision = "D"
+    """Detect the highest precision for a list of datetime objects.
+
+    Args:
+        datetime_list: List of datetime objects.
+
+    Returns:
+        A string representing the highest precision among the datetime objects
+        ('D', 's', 'ms', 'us', 'ns').
+    """
+    # Define precision hierarchy
+    precision_hierarchy = ["D", "s", "ms", "us", "ns"]
+    highest_precision_index = 0  # Start with the lowest precision ("D")

     for dt in datetime_list:
-        if dt.microsecond != 0 and dt.microsecond % 1000 != 0:
-            highest_precision = "us"
+        # Safely get the nanosecond value using getattr for backward compatibility
+        nanosecond = getattr(dt, "nanosecond", 0)
+        if nanosecond != 0:
+            current_precision = "ns"
+        elif dt.microsecond != 0:
+            # Check if the microsecond precision is exactly millisecond
+            if dt.microsecond % 1000 == 0:
+                current_precision = "ms"
+            else:
+                current_precision = "us"
+        elif dt.second != 0 or dt.minute != 0 or dt.hour != 0:
+            # pyarrow does not support h or m, use s for those cases to
+            current_precision = "s"
+        else:
+            current_precision = "D"
+
+        # Update highest_precision_index based on the hierarchy
+        current_index = precision_hierarchy.index(current_precision)
+        highest_precision_index = max(highest_precision_index, current_index)
+
+        # Stop early if highest possible precision is reached
+        if highest_precision_index == len(precision_hierarchy) - 1:
             break
-        elif dt.microsecond != 0 and dt.microsecond % 1000 == 0:
-            highest_precision = "ms"
-        elif dt.hour != 0 or dt.minute != 0 or dt.second != 0:
-            # pyarrow does not support h or m, use s for those cases too
-            highest_precision = "s"

-    return highest_precision
+    return precision_hierarchy[highest_precision_index]
+
+
+def _convert_to_datetime64(dt: datetime, precision: str) -> np.datetime64:
+    """
+    Converts a datetime object to a numpy datetime64 object with the specified
+    precision.
+
+    Args:
+        dt: A datetime object to be converted.
+        precision: The desired precision for the datetime64 conversion. Possible
+        values are 'D', 's', 'ms', 'us', 'ns'.
+
+    Returns:
+        np.datetime64: A numpy datetime64 object with the specified precision.
+    """
+    if precision == "ns":
+        # Calculate nanoseconds from microsecond and nanosecond
+        microseconds_as_ns = dt.microsecond * 1000
+        # Use getattr for backward compatibility where nanosecond attribute may not
+        # exist
+        nanoseconds = getattr(dt, "nanosecond", 0)
+        total_nanoseconds = microseconds_as_ns + nanoseconds
+        # Create datetime64 from base datetime with microsecond precision
+        base_dt = np.datetime64(dt, "us")
+        # Add remaining nanoseconds as timedelta
+        return base_dt + np.timedelta64(total_nanoseconds - microseconds_as_ns, "ns")
+    else:
+        return np.datetime64(dt).astype(f"datetime64[{precision}]")

 def _convert_datetime_list_to_array(datetime_list: List[datetime]) -> np.ndarray:
+    """Convert a list of datetime objects to a NumPy array of datetime64 with proper
+    precision.
+
+    Args:
+        datetime_list (List[datetime]): A list of `datetime` objects to be converted.
+        Each `datetime` object represents a specific point in time.
+
+    Returns:
+        np.ndarray: A NumPy array containing the `datetime64` values of the datetime
+        objects from the input list, with the appropriate precision (e.g., nanoseconds,
+        microseconds, milliseconds, etc.).
+    """
+    # Detect the highest precision for the datetime objects
     precision = _detect_highest_datetime_precision(datetime_list)

-    return np.array(
-        [np.datetime64(dt, precision) for dt in datetime_list],
-        dtype=f"datetime64[{precision}]",
-    )
+    # Convert each datetime to the corresponding numpy datetime64 with the appropriate
+    # precision
+    return np.array([_convert_to_datetime64(dt, precision) for dt in datetime_list])

 def convert_to_numpy(column_values: Any) -> np.ndarray:

diff --git a/python/ray/data/tests/test_arrow_block.py b/python/ray/data/tests/test_arrow_block.py
index ec6342c7e2b4c..80c916a22b36b 100644
--- a/python/ray/data/tests/test_arrow_block.py
+++ b/python/ray/data/tests/test_arrow_block.py
@@ -6,6 +6,7 @@
 from typing import Union

 import numpy as np
+import pandas as pd
 import pyarrow as pa
 import pytest
 from pyarrow import parquet as pq
@@ -355,6 +356,32 @@ def test_build_block_with_null_column(ray_start_regular_shared):
     assert np.array_equal(rows[1]["array"], np.zeros((2, 2)))

+def test_arrow_block_timestamp_ns(ray_start_regular_shared):
+    # Input data with nanosecond precision timestamps
+    data_rows = [
+        {"col1": 1, "col2": pd.Timestamp("2023-01-01T00:00:00.123456789")},
+        {"col1": 2, "col2": pd.Timestamp("2023-01-01T01:15:30.987654321")},
+        {"col1": 3, "col2": pd.Timestamp("2023-01-01T02:30:15.111111111")},
+        {"col1": 4, "col2": pd.Timestamp("2023-01-01T03:45:45.222222222")},
+        {"col1": 5, "col2": pd.Timestamp("2023-01-01T05:00:00.333333333")},
+    ]
+
+    # Initialize ArrowBlockBuilder
+    arrow_builder = ArrowBlockBuilder()
+    for row in data_rows:
+        arrow_builder.add(row)
+    arrow_block = arrow_builder.build()
+
+    assert arrow_block.schema.field("col2").type == pa.timestamp("ns")
+    for i, row in enumerate(data_rows):
+        result_timestamp = arrow_block["col2"][i].as_py()
+        # Convert both values to pandas Timestamp to preserve nanosecond precision for
+        # comparison.
+        assert pd.Timestamp(row["col2"]) == pd.Timestamp(
+            result_timestamp
+        ), f"Timestamp mismatch at row {i} in ArrowBlockBuilder output"
+
+
 def test_arrow_nan_element():
     ds = ray.data.from_items(
         [
diff --git a/python/ray/data/tests/test_map.py b/python/ray/data/tests/test_map.py
index 5df6ff3066d2f..8321864801f83 100644
--- a/python/ray/data/tests/test_map.py
+++ b/python/ray/data/tests/test_map.py
@@ -330,6 +330,144 @@ def map_generator(item: dict) -> Iterator[int]:
     ]

+# Helper function to process timestamp data in nanoseconds
+def process_timestamp_data(row):
+    # Convert numpy.datetime64 to pd.Timestamp if needed
+    if isinstance(row["timestamp"], np.datetime64):
+        row["timestamp"] = pd.Timestamp(row["timestamp"])
+
+    # Add 1ns to timestamp
+    row["timestamp"] = row["timestamp"] + pd.Timedelta(1, "ns")
+
+    # Ensure the timestamp column is in the expected dtype (datetime64[ns])
+    row["timestamp"] = pd.to_datetime(row["timestamp"], errors="raise")
+
+    return row
+
+
+def process_timestamp_data_batch_arrow(batch: pa.Table) -> pa.Table:
+    # Convert pyarrow Table to pandas DataFrame to process the timestamp column
+    df = batch.to_pandas()
+
+    df["timestamp"] = df["timestamp"].apply(
+        lambda x: pd.Timestamp(x) if isinstance(x, np.datetime64) else x
+    )
+
+    # Add 1ns to timestamp
+    df["timestamp"] = df["timestamp"] + pd.Timedelta(1, "ns")
+
+    # Convert back to pyarrow Table
+    return pa.table(df)
+
+
+def process_timestamp_data_batch_pandas(batch: pd.DataFrame) -> pd.DataFrame:
+    # Add 1ns to timestamp column
+    batch["timestamp"] = batch["timestamp"] + pd.Timedelta(1, "ns")
+    return batch
+
+
+@pytest.mark.parametrize(
+    "df, expected_df",
+    [
+        pytest.param(
+            pd.DataFrame(
+                {
+                    "id": [1, 2, 3],
+                    "timestamp": pd.to_datetime(
+                        [
+                            "2024-01-01 00:00:00.123456789",
+                            "2024-01-02 00:00:00.987654321",
+                            "2024-01-03 00:00:00.111222333",
+                        ]
+                    ),
+                    "value": [10.123456789, 20.987654321, 30.111222333],
+                }
+            ),
+            pd.DataFrame(
+                {
+                    "id": [1, 2, 3],
+                    "timestamp": pd.to_datetime(
+                        [
+                            "2024-01-01 00:00:00.123456790",
+                            "2024-01-02 00:00:00.987654322",
+                            "2024-01-03 00:00:00.111222334",
+                        ]
+                    ),
+                    "value": [10.123456789, 20.987654321, 30.111222333],
+                }
+            ),
+            id="nanoseconds_increment",
+        )
+    ],
+)
+def test_map_batches_timestamp_nanosecs(df, expected_df, ray_start_regular_shared):
+    """Verify handling timestamp with nanosecs in map_batches"""
+    ray_data = ray.data.from_pandas(df)
+
+    # Using pyarrow format
+    result_arrow = ray_data.map_batches(
+        process_timestamp_data_batch_arrow, batch_format="pyarrow"
+    )
+    processed_df_arrow = result_arrow.to_pandas()
+    processed_df_arrow["timestamp"] = processed_df_arrow["timestamp"].astype(
+        "datetime64[ns]"
+    )
+    pd.testing.assert_frame_equal(processed_df_arrow, expected_df)
+
+    # Using pandas format
+    result_pandas = ray_data.map_batches(
+        process_timestamp_data_batch_pandas, batch_format="pandas"
+    )
+    processed_df_pandas = result_pandas.to_pandas()
+    processed_df_pandas["timestamp"] = processed_df_pandas["timestamp"].astype(
+        "datetime64[ns]"
+    )
+    pd.testing.assert_frame_equal(processed_df_pandas, expected_df)
+
+
+@pytest.mark.parametrize(
+    "df, expected_df",
+    [
+        pytest.param(
+            pd.DataFrame(
+                {
+                    "id": [1, 2, 3],
+                    "timestamp": pd.to_datetime(
+                        [
+                            "2024-01-01 00:00:00.123456789",
+                            "2024-01-02 00:00:00.987654321",
+                            "2024-01-03 00:00:00.111222333",
+                        ]
+                    ),
+                    "value": [10.123456789, 20.987654321, 30.111222333],
+                }
+            ),
+            pd.DataFrame(
+                {
+                    "id": [1, 2, 3],
+                    "timestamp": pd.to_datetime(
+                        [
+                            "2024-01-01 00:00:00.123456790",
+                            "2024-01-02 00:00:00.987654322",
+                            "2024-01-03 00:00:00.111222334",
+                        ]
+                    ),
+                    "value": [10.123456789, 20.987654321, 30.111222333],
+                }
+            ),
+            id="nanoseconds_increment_map",
+        )
+    ],
+)
+def test_map_timestamp_nanosecs(df, expected_df, ray_start_regular_shared):
+    """Verify handling timestamp with nanosecs in map"""
+    ray_data = ray.data.from_pandas(df)
+    result = ray_data.map(process_timestamp_data)
+    processed_df = result.to_pandas()
+    processed_df["timestamp"] = processed_df["timestamp"].astype("datetime64[ns]")
+    pd.testing.assert_frame_equal(processed_df, expected_df)
+
+
 def test_add_column(ray_start_regular_shared):
     """Tests the add column API."""

diff --git a/python/ray/data/tests/test_numpy_support.py b/python/ray/data/tests/test_numpy_support.py
index ec67bcf689bb8..adca6e4ab5aa4 100644
--- a/python/ray/data/tests/test_numpy_support.py
+++ b/python/ray/data/tests/test_numpy_support.py
@@ -1,6 +1,7 @@
 from datetime import datetime

 import numpy as np
+import pandas as pd
 import pytest
 import torch

@@ -58,6 +59,7 @@ def test_list_of_objects(ray_start_regular_shared, restore_data_context):
     assert_structure_equals(output, np.array([1, 2, 3, UserObj()]))

+# Define datetime values with different precisions
 DATETIME_DAY_PRECISION = datetime(year=2024, month=1, day=1)
 DATETIME_HOUR_PRECISION = datetime(year=2024, month=1, day=1, hour=1)
 DATETIME_MIN_PRECISION = datetime(year=2024, month=1, day=1, minute=1)
@@ -65,14 +67,28 @@ def test_list_of_objects(ray_start_regular_shared, restore_data_context):
 DATETIME_MILLISEC_PRECISION = datetime(year=2024, month=1, day=1, microsecond=1000)
 DATETIME_MICROSEC_PRECISION = datetime(year=2024, month=1, day=1, microsecond=1)

+# Define pandas values for different precisions
+PANDAS_DAY_PRECISION = pd.Timestamp(year=2024, month=1, day=1)
+PANDAS_HOUR_PRECISION = pd.Timestamp(year=2024, month=1, day=1, hour=1)
+PANDAS_MIN_PRECISION = pd.Timestamp(year=2024, month=1, day=1, minute=1)
+PANDAS_SEC_PRECISION = pd.Timestamp(year=2024, month=1, day=1, second=1)
+PANDAS_MILLISEC_PRECISION = pd.Timestamp(year=2024, month=1, day=1, microsecond=1000)
+PANDAS_MICROSEC_PRECISION = pd.Timestamp(year=2024, month=1, day=1, microsecond=1)
+PANDAS_NANOSEC_PRECISION = pd.Timestamp(
+    year=2024, month=1, day=1, hour=1, minute=1, second=1
+) + pd.Timedelta(nanoseconds=1)
+
+# Define numpy.datetime64 values for comparison
 DATETIME64_DAY_PRECISION = np.datetime64("2024-01-01")
 DATETIME64_HOUR_PRECISION = np.datetime64("2024-01-01T01:00", "s")
 DATETIME64_MIN_PRECISION = np.datetime64("2024-01-01T00:01", "s")
 DATETIME64_SEC_PRECISION = np.datetime64("2024-01-01T00:00:01")
 DATETIME64_MILLISEC_PRECISION = np.datetime64("2024-01-01T00:00:00.001")
 DATETIME64_MICROSEC_PRECISION = np.datetime64("2024-01-01T00:00:00.000001")
+DATETIME64_NANOSEC_PRECISION = np.datetime64("2024-01-01T01:01:01.000000001", "ns")

+# Parametrized test to validate datetime values and expected numpy.datetime64 results
 @pytest.mark.parametrize(
     "data,expected_output",
     [
@@ -85,17 +101,59 @@ def test_list_of_objects(ray_start_regular_shared, restore_data_context):
         (
             [DATETIME_MICROSEC_PRECISION, DATETIME_MILLISEC_PRECISION],
             np.array(
-                [DATETIME64_MICROSEC_PRECISION, DATETIME_MILLISEC_PRECISION],
+                [DATETIME64_MICROSEC_PRECISION, DATETIME64_MILLISEC_PRECISION],
                 dtype="datetime64[us]",
             ),
         ),
         (
             [DATETIME_SEC_PRECISION, DATETIME_MILLISEC_PRECISION],
             np.array(
-                [DATETIME64_SEC_PRECISION, DATETIME_MILLISEC_PRECISION],
+                [DATETIME64_SEC_PRECISION, DATETIME64_MILLISEC_PRECISION],
                 dtype="datetime64[ms]",
             ),
         ),
+        (
+            [DATETIME_DAY_PRECISION, DATETIME_SEC_PRECISION],
+            np.array(
+                [DATETIME64_DAY_PRECISION, DATETIME64_SEC_PRECISION],
+                dtype="datetime64[s]",
+            ),
+        ),
+        ([PANDAS_DAY_PRECISION], np.array([DATETIME64_DAY_PRECISION])),
+        ([PANDAS_HOUR_PRECISION], np.array([DATETIME64_HOUR_PRECISION])),
+        ([PANDAS_MIN_PRECISION], np.array([DATETIME64_MIN_PRECISION])),
+        ([PANDAS_SEC_PRECISION], np.array([DATETIME64_SEC_PRECISION])),
+        ([PANDAS_MILLISEC_PRECISION], np.array([DATETIME64_MILLISEC_PRECISION])),
+        ([PANDAS_MICROSEC_PRECISION], np.array([DATETIME64_MICROSEC_PRECISION])),
+        ([PANDAS_NANOSEC_PRECISION], np.array([DATETIME64_NANOSEC_PRECISION])),
+        (
+            [PANDAS_NANOSEC_PRECISION, PANDAS_MICROSEC_PRECISION],
+            np.array(
+                [DATETIME64_NANOSEC_PRECISION, DATETIME64_MICROSEC_PRECISION],
+                dtype="datetime64[ns]",
+            ),
+        ),
+        (
+            [PANDAS_MICROSEC_PRECISION, PANDAS_MILLISEC_PRECISION],
+            np.array(
+                [DATETIME64_MICROSEC_PRECISION, DATETIME64_MILLISEC_PRECISION],
+                dtype="datetime64[us]",
+            ),
+        ),
+        (
+            [PANDAS_SEC_PRECISION, PANDAS_MILLISEC_PRECISION],
+            np.array(
+                [DATETIME64_SEC_PRECISION, DATETIME64_MILLISEC_PRECISION],
+                dtype="datetime64[ms]",
+            ),
+        ),
+        (
+            [PANDAS_DAY_PRECISION, PANDAS_SEC_PRECISION],
+            np.array(
+                [DATETIME64_DAY_PRECISION, DATETIME64_SEC_PRECISION],
+                dtype="datetime64[s]",
+            ),
+        ),
     ],
 )
 def test_list_of_datetimes(
diff --git a/python/ray/data/tests/test_pandas_block.py b/python/ray/data/tests/test_pandas_block.py
index f7e98f4bd1eb1..4e2e1a93c9d40 100644
--- a/python/ray/data/tests/test_pandas_block.py
+++ b/python/ray/data/tests/test_pandas_block.py
@@ -9,7 +9,7 @@

 import ray
 import ray.data
-from ray.data._internal.pandas_block import PandasBlockAccessor
+from ray.data._internal.pandas_block import PandasBlockAccessor, PandasBlockBuilder
 from ray.data.extensions.object_extension import _object_extension_type_allowed

 # Set seed for the test for size as it related to sampling
@@ -28,6 +28,32 @@ def test_append_column(ray_start_regular_shared):
     assert actual_block.equals(expected_block)

+def test_pandas_block_timestamp_ns(ray_start_regular_shared):
+    # Input data with nanosecond precision timestamps
+    data_rows = [
+        {"col1": 1, "col2": pd.Timestamp("2023-01-01T00:00:00.123456789")},
+        {"col1": 2, "col2": pd.Timestamp("2023-01-01T01:15:30.987654321")},
+        {"col1": 3, "col2": pd.Timestamp("2023-01-01T02:30:15.111111111")},
+        {"col1": 4, "col2": pd.Timestamp("2023-01-01T03:45:45.222222222")},
+        {"col1": 5, "col2": pd.Timestamp("2023-01-01T05:00:00.333333333")},
+    ]
+
+    # Initialize PandasBlockBuilder
+    pandas_builder = PandasBlockBuilder()
+    for row in data_rows:
+        pandas_builder.add(row)
+    pandas_block = pandas_builder.build()
+
+    assert pd.api.types.is_datetime64_ns_dtype(pandas_block["col2"])
+
+    for original_row, result_row in zip(
+        data_rows, pandas_block.to_dict(orient="records")
+    ):
+        assert (
+            original_row["col2"] == result_row["col2"]
+        ), "Timestamp mismatch in PandasBlockBuilder output"
+
+
 @pytest.mark.skipif(
     _object_extension_type_allowed(), reason="Objects can be put into Arrow"
 )

[Data] Inconsistent behavior with Ray Data and timestamps
### What happened + What you expected to happen

*This example has been provided as a gist instead of in-line code to make it easier to reproduce and trace.

Two odd behaviors are identified:
1. When working with flat-structured datasets, ray data performs destructive type conversions when using the `.map()` API.
2. When working with nest-structed datasets (ie. data with sub-structs), ray data performs destructive type conversions when using the `.map_batches()` (or `.take_batches()`) API.

The first issue is isolated to just `ray.data`, while the second is demonstrated to fail with the `pandas` API as well.

This currently blocks me from using ray data as I cannot extract timestamps consistently across my data sources.

### Versions / Dependencies

Name: numpy
Version: 1.24.4

Name: pandas
Version: 2.0.3

Name: pyarrow
Version: 17.0.0

Name: ray
Version: 2.10.0

Name: host
Version: Darwin [hostname-redacted] 23.5.0 Darwin Kernel Version 23.5.0: Wed May  1 20:13:18 PDT 2024; root:xnu-10063.121.3~5/RELEASE_ARM64_T6030 arm64

### Reproduction script

Gist: https://gist.github.com/NellyWhads/fdfb261a027be7e7bc87bec91d9e9035

### Issue Severity

High: It blocks me from completing my task.

@jcotant1 the crux for both issues is around how we handle nanoseconds

1. for the 1st behavior: `.map()` + `.take_batch(..., batch_format="pandas")` truncates the time information (see the **ns** field)
2. for the 2nd behavior (timestamps within a struct):
    - currently it casts timestamps with nanoseconds to integers while those with lower precisions are preserved as Python datetime objects. This prohibits any caller from relying on the type, and requires them to coerce the value back to the appropriate type before every use
    - since Python's datetime and timedelta objects don't support nanoseconds, @NellyWhads asked if Ray can standardize around np.datetime64 and/or pd.Timestamp (perhaps depending on the batch format)?
Based on the warnings logged from ray data, `map_batch()` / `take_batch()` is the API users should target since it is more consistent with the interface to the ray training runner. If this is the case, the 2nd issue is of much higher priority than the first.
Test case for this.

```
import ray
import pandas as pd

# Initialize Ray
ray.init(ignore_reinit_error=True)

# Sample flat-structured dataset (pandas DataFrame)
df = pd.DataFrame({
    'id': [1, 2, 3],
    'timestamp': pd.to_datetime(['2024-01-01 00:00:00.123456789',
                                 '2024-01-02 00:00:00.987654321',
                                 '2024-01-03 00:00:00.111222333']),
    'value': [10.123456789, 20.987654321, 30.111222333]
})

# Ray data setup
ray_data = ray.data.from_pandas(df)

# Define a function to apply with .map()
def process_data(row):
    # Example of applying a transformation
    row['timestamp'] = row['timestamp'] + pd.Timedelta(1, 'ns')  # add 1ns to timestamp
    return row

# Apply the .map() function
result = ray_data.map(process_data)

# Convert back to pandas for inspection
processed_df = result.to_pandas()

# Print the original and processed data
print("Original DataFrame:")
print(df)
print("\nProcessed DataFrame after .map():")
print(processed_df)

```

```
Original DataFrame:
   id                     timestamp      value
0   1 2024-01-01 00:00:00.123456789  10.123457
1   2 2024-01-02 00:00:00.987654321  20.987654
2   3 2024-01-03 00:00:00.111222333  30.111222

Processed DataFrame after .map():
   id                  timestamp      value
0   1 2024-01-01 00:00:00.123456  10.123457
1   2 2024-01-02 00:00:00.987654  20.987654
2   3 2024-01-03 00:00:00.111222  30.111222
```
for the 2nd issue, i believe it's because Ray Data relies on PyArrow's Pandas integration. PyArrow doesn't seem to respect the conversion table [here](https://arrow.apache.org/docs/python/pandas.html#arrow-pandas-conversion) if the timestamps are inside a StructArray

Repro:
```
import pandas as pd
import pyarrow as pa

base_timestamp = pd.Timestamp("2024-01-01")

start_times = pa.array([base_timestamp + pd.Timedelta(nanoseconds=i) for i in range(0, 5)], type=pa.timestamp('ns'))
end_times = pa.array([base_timestamp + pd.Timedelta(nanoseconds=j) for j in range(1, 6)], type=pa.timestamp('ns'))

struct_array = pa.StructArray.from_arrays(
    [start_times, end_times],
    fields=[
        pa.field("start_time", pa.timestamp('ns')),
        pa.field("end_time", pa.timestamp('ns'))
        ]
    )
# when not using StructArray, conversion works as expected per https://arrow.apache.org/docs/python/pandas.html#arrow-pandas-conversion
table_without_struct = pa.Table.from_arrays([start_times, end_times], ["start_timestamp", "end_timestamp"])
df_without_struct = table_without_struct.to_pandas()
print(df_without_struct.iloc[0]) # returns pd.Series with dtype=datetime64[ns]

print("--------")
# however when using pyarrow StructArray, pyarrow converts timestamps with nanosecond precision to ints rather than respecting the conversion table
# this is true even without using Ray
table_with_struct = pa.Table.from_arrays([struct_array], ["timestamps"])
df_with_struct = table_with_struct.to_pandas()
print(df_with_struct["timestamps"].iloc[0]) # returns a dict of ints
I believe that's the exact issue. I'm unfortunately not sure if pandas will be able to iterate on a fix quickly in this case. Perhaps the ray team has a better relationship with the pandas devs than I (a relative stranger) may have?

A pandas-centric fix also means that users are forced to use a specific version of pandas in their environments.
Would you be able to convert the timestamp to int64? This should handle nanosecs precision.

```
import ray
import pandas as pd

# Initialize Ray
ray.init(ignore_reinit_error=True)

# Sample flat-structured dataset (pandas DataFrame)
df = pd.DataFrame({
    'id': [1, 2, 3],
    'timestamp': pd.to_datetime(['2024-01-01 00:00:00.123456789',
                                 '2024-01-02 00:00:00.987654321',
                                 '2024-01-03 00:00:00.111222333']),
    'value': [10.123456789, 20.987654321, 30.111222333]
})

# Convert timestamp to nanoseconds (int64) before passing to Ray
df['timestamp_ns'] = df['timestamp'].astype('int64')  # Convert to nanoseconds since epoch
df = df.drop(columns=['timestamp'])  # Drop original timestamp if not needed

# Ray data setup
ray_data = ray.data.from_pandas(df)

# Define a function to apply with .map()
def process_data(row):
    # Add 1 nanosecond
    row['timestamp_ns'] += 1
    return row

# Apply the .map() function
result = ray_data.map(process_data)

# Convert back to pandas for inspection
processed_df = result.to_pandas()

# Restore timestamp from nanoseconds
processed_df['timestamp'] = pd.to_datetime(processed_df['timestamp_ns'])  # Convert back to timestamp
processed_df = processed_df.drop(columns=['timestamp_ns'])  # Drop intermediate column if not needed

# Print the original and processed DataFrames
print("Original DataFrame:")
print(df)
print("\nProcessed DataFrame after .map():")
print(processed_df)
```

```
Original DataFrame:
   id      value         timestamp_ns
0   1  10.123457  1704067200123456789
1   2  20.987654  1704153600987654321
2   3  30.111222  1704240000111222333

Processed DataFrame after .map():
   id      value                     timestamp
0   1  10.123457 2024-01-01 00:00:00.123456790
1   2  20.987654 2024-01-02 00:00:00.987654322
2   3  30.111222 2024-01-03 00:00:00.111222334
```
Unfortunately, this means that whoever uses the data is now responsible for remembering to convert it back to a timestamp (in each transform they `map_batch()` onto the data).

This is how we currently "patch" the behavior, but it requires an extra layer of API that is rather expensive to maintain instead of just using ray data as a product directly.
I actually run into the same bug, however, I only became aware when ray data raised an error during `pa.concat_tables` while unifying the schemas of two blocks.

This happens in cases there the type of the timestamp columns of two blocks are inferred differently:
`2024-01-01 00:00:00.123456790 -> 2024-01-01 00:00:00.123456 -> timestamp[us] (microseconds)`, however,
`2024-01-01 00:00:00.123000090 -> 2024-01-01 00:00:00.123 -> timestamp[ms] (milliseconds)`


diff --git a/python/ray/dag/compiled_dag_node.py b/python/ray/dag/compiled_dag_node.py
index 33abb86ae734f..6e3e1449eee06 100644
--- a/python/ray/dag/compiled_dag_node.py
+++ b/python/ray/dag/compiled_dag_node.py
@@ -823,6 +823,8 @@ def __init__(
         # Preprocessing identifies the input node and output node.
         self.input_task_idx: Optional[int] = None
         self.output_task_idx: Optional[int] = None
+        # List of task indices that are input attribute nodes.
+        self.input_attr_task_idxs: List[int] = []
         # Denotes whether execute/execute_async returns a list of refs/futures.
         self._returns_list: bool = False
         # Number of expected positional args and kwargs that may be passed to
@@ -949,11 +951,13 @@ def _preprocess(self) -> None:
         nccl_actors_p2p: Set["ray.actor.ActorHandle"] = set()
         nccl_collective_ops: Set[_CollectiveOperation] = set()

-        # Find the input node to the DAG.
+        # Find the input node and input attribute nodes in the DAG.
         for idx, task in self.idx_to_task.items():
             if isinstance(task.dag_node, InputNode):
                 assert self.input_task_idx is None, "More than one InputNode found"
                 self.input_task_idx = idx
+            elif isinstance(task.dag_node, InputAttributeNode):
+                self.input_attr_task_idxs.append(idx)

         # Find the (multi-)output node to the DAG.
         for idx, task in self.idx_to_task.items():
@@ -1088,6 +1092,9 @@ def _preprocess(self) -> None:
                 ):
                     downstream_actor_handle = dag_node._get_actor_handle()

+                # Add the type hint of the upstream node to the task.
+                task.arg_type_hints.append(upstream_task.dag_node.type_hint)
+
                 if isinstance(upstream_task.dag_node, InputAttributeNode):
                     # Record all of the keys used to index the InputNode.
                     # During execution, we will check that the user provides
@@ -1123,7 +1130,6 @@ def _preprocess(self) -> None:
                     direct_input = True

                 upstream_task.downstream_task_idxs[task_idx] = downstream_actor_handle
-                task.arg_type_hints.append(upstream_task.dag_node.type_hint)

                 if upstream_task.dag_node.type_hint.requires_nccl():
                     # Add all readers to the NCCL actors of P2P.
@@ -1496,8 +1502,6 @@ def _get_or_compile(
             )

         input_task = self.idx_to_task[self.input_task_idx]
-        # Register custom serializers for inputs provided to dag.execute().
-        input_task.dag_node.type_hint.register_custom_serializer()
         self.dag_input_channels = input_task.output_channels
         assert self.dag_input_channels is not None

@@ -1599,8 +1603,9 @@ def _get_or_compile(
             task = self.idx_to_task[output_idx]
             assert len(task.output_channels) == 1
             self.dag_output_channels.append(task.output_channels[0])
-            # Register custom serializers for DAG outputs.
-            output.type_hint.register_custom_serializer()
+
+        # Register custom serializers for input, input attribute, and output nodes.
+        self._register_input_output_custom_serializer()

         assert self.dag_input_channels
         assert self.dag_output_channels
@@ -2746,6 +2751,26 @@ def visualize(
         dot.render(filename, view=view)
         return dot.source

+    def _register_input_output_custom_serializer(self):
+        """
+        Register custom serializers for input, input attribute, and output nodes.
+        """
+        assert self.input_task_idx is not None
+        assert self.output_task_idx is not None
+
+        # Register custom serializers for input node.
+        input_task = self.idx_to_task[self.input_task_idx]
+        input_task.dag_node.type_hint.register_custom_serializer()
+
+        # Register custom serializers for input attribute nodes.
+        for input_attr_task_idx in self.input_attr_task_idxs:
+            input_attr_task = self.idx_to_task[input_attr_task_idx]
+            input_attr_task.dag_node.type_hint.register_custom_serializer()
+
+        # Register custom serializers for output nodes.
+        for output in self.idx_to_task[self.output_task_idx].args:
+            output.type_hint.register_custom_serializer()
+
     def teardown(self, kill_actors: bool = False):
         """Teardown and cancel all actor tasks for this DAG. After this
         function returns, the actors should be available to execute new tasks

diff --git a/python/ray/dag/tests/experimental/test_torch_tensor_dag.py b/python/ray/dag/tests/experimental/test_torch_tensor_dag.py
index 61684af300cb8..6219ac74dacc4 100644
--- a/python/ray/dag/tests/experimental/test_torch_tensor_dag.py
+++ b/python/ray/dag/tests/experimental/test_torch_tensor_dag.py
@@ -118,6 +118,20 @@ def forward(self, inp):
         return torch.randn(10, 10)

+@ray.remote
+class Worker:
+    def __init__(self):
+        self.device = None
+
+    def echo(self, tensor):
+        assert isinstance(tensor, torch.Tensor)
+        self.device = tensor.device
+        return tensor
+
+    def get_device(self):
+        return self.device
+
+
 @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
 def test_torch_tensor_p2p(ray_start_regular):
     if USE_GPU:
@@ -1268,6 +1282,198 @@ def recv(self, tensor):
     compiled_dag.teardown()

+class TestTorchTensorTypeHintCustomSerializer:
+    # All tests inside this file are running in the same process, so we need to
+    # manually deregister the custom serializer for `torch.Tensor` before and
+    # after each test to avoid side effects.
+    def setup_method(self):
+        ray.util.serialization.deregister_serializer(torch.Tensor)
+
+    def teardown_method(self):
+        ray.util.serialization.deregister_serializer(torch.Tensor)
+
+    @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
+    @pytest.mark.parametrize("tensor_device", ["cpu", "cuda"])
+    def test_input_node_without_type_hint(self, ray_start_regular, tensor_device):
+        """
+        Since no TorchTensorType hint is provided in this compiled graph,
+        normal serialization and deserialization functions are used, which will
+        not move the tensor to GPU/CPU.
+        """
+        if not USE_GPU:
+            pytest.skip("Test requires GPU")
+
+        worker = Worker.options(num_gpus=1).remote()
+
+        with InputNode() as inp:
+            dag = worker.echo.bind(inp)
+
+        compiled_dag = dag.experimental_compile()
+        tensor = torch.tensor([5])
+        if tensor_device == "cuda":
+            tensor = tensor.cuda()
+        ref = compiled_dag.execute(tensor)
+        t = ray.get(ref)
+        assert torch.equal(t, tensor)
+
+        device = ray.get(worker.get_device.remote())
+        assert device.type == tensor_device
+
+    @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
+    @pytest.mark.parametrize("tensor_device", ["cpu", "cuda"])
+    def test_input_node_with_type_hint(self, ray_start_regular, tensor_device):
+        """
+        Since `inp` has a TorchTensorType hint, both the driver and `worker` will
+        use the custom serializer.
+
+        Step 1: The driver calls `serialize_tensor` to serialize `input_tensor` and
+                move the tensor to CPU if it is on GPU.
+        Step 2: The `worker` calls `deserialize_tensor` to deserialize `input_tensor`
+               and moves it to GPU.
+        Step 3: The `worker` calls `serialize_tensor` to serialize the result of
+               `echo` and moves it to CPU.
+        Step 4: The driver calls `deserialize_tensor` to deserialize the result of
+               `echo`. Since the driver's `ChannelContext.torch_device` is CPU,
+               the tensor will not be moved to GPU.
+        """
+        if not USE_GPU:
+            pytest.skip("Test requires GPU")
+
+        worker = Worker.options(num_gpus=1).remote()
+
+        with InputNode() as inp:
+            dag = worker.echo.bind(inp.with_type_hint(TorchTensorType()))
+        compiled_dag = dag.experimental_compile()
+        cpu_tensor = torch.tensor([1])
+        input_tensor = cpu_tensor
+        if tensor_device == "cuda":
+            input_tensor = input_tensor.cuda()
+        ref = compiled_dag.execute(input_tensor)
+        # Verify Step 4
+        t = ray.get(ref)
+        assert torch.equal(t, cpu_tensor)
+
+        # Verify Step 2
+        device = ray.get(worker.get_device.remote())
+        assert device.type == "cuda"
+
+    @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
+    def test_input_attr_nodes_with_all_tensor_type_hint(self, ray_start_regular):
+        """
+        Since both `inp[0]` and `inp[1]` have tensor type hint, both workers will
+        use the custom serializer.
+
+        Step 1: The driver calls `serialize_tensor` to serialize `cpu_tensor_1`
+        and `cpu_tensor_2`.
+
+        Step 2:
+        * The `worker1` calls `deserialize_tensor` to deserialize `cpu_tensor_1`
+          and moves it to GPU.
+        * The `worker2` calls `deserialize_tensor` to deserialize `cpu_tensor_2`
+          and moves it to GPU.
+
+        Step 3:
+        * The `worker1` calls `serialize_tensor` to serialize the result of
+          `echo` and moves it to CPU.
+        * The `worker2` calls `serialize_tensor` to serialize the result of
+          `echo` and moves it to CPU.
+
+        Step 4: The driver calls `deserialize_tensor` to deserialize the result
+        of `echo`. Since the driver's `ChannelContext.torch_device` is CPU,
+        the tensor will not be moved to GPU.
+        """
+        if not USE_GPU:
+            pytest.skip("Test requires GPU")
+
+        worker1 = Worker.options(num_gpus=1).remote()
+        worker2 = Worker.options(num_gpus=1).remote()
+        with InputNode() as inp:
+            dag = inp[0].with_type_hint(TorchTensorType())
+            branch1 = worker1.echo.bind(dag)
+            dag = inp[1].with_type_hint(TorchTensorType())
+            branch2 = worker2.echo.bind(dag)
+            dag = MultiOutputNode([branch1, branch2])
+
+        compiled_dag = dag.experimental_compile()
+        cpu_tensor_1 = torch.tensor([1])
+        cpu_tensor_2 = torch.tensor([2])
+        ref = compiled_dag.execute(cpu_tensor_1, cpu_tensor_2)
+
+        # Verify Step 4
+        t1, t2 = ray.get(ref)
+        assert torch.equal(t1, cpu_tensor_1)
+        assert torch.equal(t2, cpu_tensor_2)
+
+        # Verify Step 2
+        device1 = ray.get(worker1.get_device.remote())
+        device2 = ray.get(worker2.get_device.remote())
+        assert device1.type == "cuda"
+        assert device2.type == "cuda"
+
+    @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
+    def test_input_attr_nodes_with_and_without_type_hint(self, ray_start_regular):
+        """
+        Only `inp[0]` has a tensor type hint, so only `worker1` will use the custom
+        serializer. Note that although we don't register the custom serializer for
+        `worker2`, it still uses the custom deserializer. This is because when custom
+        serializers are registered with Ray, the registered deserializer is shipped
+        with the serialized value and used on the receiving end. See the comment in
+        `ChannelOutputType.register_custom_serializer` for more details.
+
+        Step 1: The driver calls `serialize_tensor` to serialize `cpu_tensor_1`
+        and `cpu_tensor_2`.
+
+        Step 2:
+        * The `worker1` calls `deserialize_tensor` to deserialize `cpu_tensor_1`
+          and moves it to GPU.
+        * The `worker2` calls `deserialize_tensor` to deserialize `cpu_tensor_2`
+          and moves it to GPU.
+
+        Step 3:
+        * The `worker1` calls `serialize_tensor` to serialize the result of `echo`
+          and moves it to CPU.
+        * The `worker2` calls the normal serialization function to serialize the
+          result of `echo` because it doesn't have a custom serializer, so the
+          tensor is still on GPU.
+
+        Step 4:
+        * The driver calls `deserialize_tensor` to deserialize the tensor from
+          `worker1`. Since the driver's `ChannelContext.torch_device` is CPU,
+          the tensor will not be moved to GPU.
+        * The driver calls normal deserialization function to deserialize the
+          tensor from `worker2`.
+        """
+        if not USE_GPU:
+            pytest.skip("Test requires GPU")
+
+        worker1 = Worker.options(num_gpus=1).remote()
+        worker2 = Worker.options(num_gpus=1).remote()
+
+        with InputNode() as inp:
+            dag = inp[0].with_type_hint(TorchTensorType())
+            branch1 = worker1.echo.bind(dag)
+            dag = inp[1]
+            branch2 = worker2.echo.bind(dag)
+            dag = MultiOutputNode([branch1, branch2])
+
+        compiled_dag = dag.experimental_compile()
+        cpu_tensor_1 = torch.tensor([1])
+        cpu_tensor_2 = torch.tensor([2])
+        ref = compiled_dag.execute(cpu_tensor_1, cpu_tensor_2)
+        t1, t2 = ray.get(ref)
+        # Verify Step 3-1
+        assert torch.equal(t1, cpu_tensor_1)
+        # Verify Step 3-2
+        gpu_tensor_2 = cpu_tensor_2.cuda()
+        assert torch.equal(t2, gpu_tensor_2)
+
+        # Verify Step 2
+        device1 = ray.get(worker1.get_device.remote())
+        device2 = ray.get(worker2.get_device.remote())
+        assert device1.type == "cuda"
+        assert device2.type == "cuda"
+
+
 @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
 def test_torch_nccl_channel_with_local_reader(ray_start_regular):
     if not USE_GPU:

[aDAG][Core] Feeding multiple inputs to aDAG and annotating with TorchTensor does not move tensors to GPU
### What happened + What you expected to happen

When feeding multiple tensors as input to an aDAG, the tensors do not get moved to GPU memory despite annotating with TorchTensor(). The same issue does not occur when only feeding one tensor to the aDAG.

### Versions / Dependencies

Ray==2.34

### Reproduction script

```python
import ray
import torch
import ray.dag
from ray.experimental.channel.torch_tensor_type import TorchTensorType

@ray.remote(num_gpus=1)
class Actor:
    def __init__(self) -> None:
        pass

    def test(self, tensor: torch.Tensor):
        return tensor.device

if __name__ == "__main__":
    ray.init()
    actor1, actor2 = Actor.remote(), Actor.remote()
    with ray.dag.InputNode() as dag_input:
        in1, in2 = dag_input[0], dag_input[1]

        # Dag 1
        in1 = in1.with_type_hint(TorchTensorType())
        in1 = actor1.test.bind(in1)

        # Dag 2
        in2 = in2.with_type_hint(TorchTensorType())
        in2 = actor2.test.bind(in2)

        dag = ray.dag.MultiOutputNode([in1, in2])

    adag = dag.experimental_compile()
    output = ray.get(adag.execute(torch.randn(2, 16), torch.tensor(1.0)))
    print(output)
```

### Issue Severity

Medium: It is a significant difficulty but I can work around it.

@stephanie-wang
cc @kevin85421 > supporting multi-channel on single input will enable ADAG to support this...


diff --git a/python/ray/data/dataset.py b/python/ray/data/dataset.py
index d2e24d39cfa8..75f98de9f2e4 100644
--- a/python/ray/data/dataset.py
+++ b/python/ray/data/dataset.py
@@ -5,7 +5,6 @@
 import logging
 import time
 import warnings
-from collections.abc import Sequence
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -765,20 +764,9 @@ def add_column(
                 f"got: {batch_format}"
             )

-        def _raise_duplicate_column_error(col: str):
-            raise ValueError(f"Trying to add an existing column with name {col!r}")
-
         def add_column(batch: DataBatch) -> DataBatch:
             column = fn(batch)
             if batch_format == "pandas":
-                import pandas as pd
-
-                assert isinstance(column, (pd.Series, Sequence)), (
-                    f"For pandas batch format, the function must return a pandas "
-                    f"Series or sequence, got: {type(column)}"
-                )
-                if col in batch:
-                    _raise_duplicate_column_error(col)
                 batch.loc[:, col] = column
                 return batch
             elif batch_format == "pyarrow":
@@ -797,10 +785,9 @@ def add_column(batch: DataBatch) -> DataBatch:
                 # which case we'll want to append it
                 column_idx = batch.schema.get_field_index(col)
                 if column_idx == -1:
-                    # Append the column to the table
                     return batch.append_column(col, column)
                 else:
-                    _raise_duplicate_column_error(col)
+                    return batch.set_column(column_idx, col, column)

             else:
                 # batch format is assumed to be numpy since we checked at the
@@ -809,8 +796,6 @@ def add_column(batch: DataBatch) -> DataBatch:
                     f"For numpy batch format, the function must return a "
                     f"numpy.ndarray, got: {type(column)}"
                 )
-                if col in batch:
-                    _raise_duplicate_column_error(col)
                 batch[col] = column
                 return batch

diff --git a/python/ray/data/tests/test_map.py b/python/ray/data/tests/test_map.py
index 0db0681d8d61..a29e5925815f 100644
--- a/python/ray/data/tests/test_map.py
+++ b/python/ray/data/tests/test_map.py
@@ -350,15 +350,11 @@ def test_add_column(ray_start_regular_shared):
     )
     assert ds.take(1) == [{"id": 0, "foo": 1}]

-    # Adding a column that is already there should result in an error
-    with pytest.raises(
-        ray.exceptions.UserCodeException,
-        match="Trying to add an existing column with name 'id'",
-    ):
-        ds = ray.data.range(5).add_column(
-            "id", lambda x: pc.add(x["id"], 1), batch_format="pyarrow"
-        )
-        assert ds.take(2) == [{"id": 1}, {"id": 2}]
+    # Adding a column that is already there should not result in an error
+    ds = ray.data.range(5).add_column(
+        "id", lambda x: pc.add(x["id"], 1), batch_format="pyarrow"
+    )
+    assert ds.take(2) == [{"id": 1}, {"id": 2}]

     # Adding a column in the wrong format should result in an error
     with pytest.raises(
@@ -378,15 +374,11 @@ def test_add_column(ray_start_regular_shared):
     )
     assert ds.take(1) == [{"id": 0, "foo": 1}]

-    # Adding a column that is already there should result in an error
-    with pytest.raises(
-        ray.exceptions.UserCodeException,
-        match="Trying to add an existing column with name 'id'",
-    ):
-        ds = ray.data.range(5).add_column(
-            "id", lambda x: np.add(x["id"], 1), batch_format="numpy"
-        )
-        assert ds.take(2) == [{"id": 1}, {"id": 2}]
+    # Adding a column that is already there should not result in an error
+    ds = ray.data.range(5).add_column(
+        "id", lambda x: np.add(x["id"], 1), batch_format="numpy"
+    )
+    assert ds.take(2) == [{"id": 1}, {"id": 2}]

     # Adding a column in the wrong format should result in an error
     with pytest.raises(
@@ -402,23 +394,20 @@ def test_add_column(ray_start_regular_shared):
     ds = ray.data.range(5).add_column("foo", lambda x: x["id"] + 1)
     assert ds.take(1) == [{"id": 0, "foo": 1}]

-    # Adding a column that is already there should result in an error
-    with pytest.raises(
-        ray.exceptions.UserCodeException,
-        match="Trying to add an existing column with name 'id'",
-    ):
-        ds = ray.data.range(5).add_column("id", lambda x: x["id"] + 1)
-        assert ds.take(2) == [{"id": 1}, {"id": 2}]
+    # Adding a column that is already there should not result in an error
+    ds = ray.data.range(5).add_column("id", lambda x: x["id"] + 1)
+    assert ds.take(2) == [{"id": 1}, {"id": 2}]

-    # Adding a column in the wrong format should result in an error
-    with pytest.raises(
-        ray.exceptions.UserCodeException, match="For pandas batch format"
-    ):
+    # Adding a column in the wrong format may result in an error
+    with pytest.raises(ray.exceptions.UserCodeException):
         ds = ray.data.range(5).add_column(
-            "id", lambda x: np.array([1]), batch_format="pandas"
+            "id", lambda x: range(7), batch_format="pandas"
         )
         assert ds.take(2) == [{"id": 1}, {"id": 2}]

+    ds = ray.data.range(5).add_column("const", lambda _: 3, batch_format="pandas")
+    assert ds.take(2) == [{"id": 0, "const": 3}, {"id": 1, "const": 3}]
+
     with pytest.raises(ValueError):
         ds = ray.data.range(5).add_column("id", 0)

[Data] Ray 2.40 `add_column` Erroneous Assertion
### What happened + What you expected to happen

Ray 2.40 introduces an assertion that a value is of type `pd.Series`, but this is erroneous. A dataframe may have a constant assigned as in the following example or be a geopandas series (which should subclass).

```python
>>> import pandas as pd
>>>
>>> # Create a sample DataFrame
...     "A": [1, 2, 3],
...     "B": [4, 5, 6]
... })
>>>
>>> # Assign a constant value to a new column "C"
>>> df["C"] = 10
>>>
>>> print(df)
   A  B   C
0  1  4  10
1  2  5  10
2  3  6  10
>>>
```

As a result of the change valid expressions like:
```python
q_ds = q_ds.add_column("_q", lambda _: True)
s_ds = s_ds.add_column("_q", lambda _: False)
```
no longer succeed.

### Versions / Dependencies

Ray 2.40

### Reproduction script
```python
import pandas as pd
import ray

df= pd.DataFrame({"geometry": [1, 2, 3]})

# Create Ray datasets
ds = ray.data.from_pandas(df)

ds = ds.add_column("works", lambda _: True)
ds = ds.add_column("works", lambda _: True)

# evaluates and fails
ds.materialize()
ds.show()
```
### Issue Severity

Low: It annoys or frustrates me.

As a workaround you should be able to do `pd.Series(constant, index=df.index)`
Here is another issue with the function update; the [docstring says columns may be overridden](https://github.com/ray-project/ray/blob/8fe05532e7037aac1b9b3755f0053047dc0f1321/python/ray/data/dataset.py#L743C13-L744C39), but there is a [check raising an error if the column already exists](https://github.com/ray-project/ray/blob/8fe05532e7037aac1b9b3755f0053047dc0f1321/python/ray/data/dataset.py#L780C17-L781C55).

As a side note, after using `add_column` you can no longer get the `ds.columns()` without the `Dataset` being at least partially evaluated, but it should be possible to know the new column names as it is the same as before but with the additional column name provided as the first argument to `add_column`. My thought being it would be good to have a function that can add a column iff the column is not present without the cost of materializing the dataset like `if col not in ds.columns(): ds.add_column(col, func)`. However, this can also be accomplished with (pseudo-code lambda) `ds.map_batches(lambda df: (df[col] = func(df)) if col not in df else df)`.
@ivanthewebber good catch -- would you be willing to submit a PR?
If someone gets to it first I am finding other good first issues as well.


diff --git a/python/ray/dag/dag_node_operation.py b/python/ray/dag/dag_node_operation.py
index ea71a88fb5e7..6073e58c72f7 100644
--- a/python/ray/dag/dag_node_operation.py
+++ b/python/ray/dag/dag_node_operation.py
@@ -435,7 +435,8 @@ def _build_dag_node_operation_graph(
             isinstance(task.dag_node, ClassMethodNode)
             and task.dag_node.is_class_method_output
         ):
-            # TODO(wxdeng): Handle the case where the task is a class method output.
+            # Class method output node dependencies are handled at its upstream:
+            # i.e., class method node
             continue
         for downstream_task_idx in task.downstream_task_idxs:
             downstream_dag_node = idx_to_task[downstream_task_idx].dag_node
@@ -445,8 +446,18 @@ def _build_dag_node_operation_graph(
                 isinstance(downstream_dag_node, ClassMethodNode)
                 and downstream_dag_node.is_class_method_output
             ):
-                # TODO(wxdeng): Handle the case where the downstream task is
-                # a class method output.
+                consumer_idxs = idx_to_task[downstream_task_idx].downstream_task_idxs
+                for consumer_idx in consumer_idxs:
+                    if consumer_idx in graph:
+                        _add_edge(
+                            graph[task_idx][_DAGNodeOperationType.WRITE],
+                            graph[consumer_idx][_DAGNodeOperationType.READ],
+                            "nccl"
+                            if graph[task_idx][
+                                _DAGNodeOperationType.WRITE
+                            ].requires_nccl
+                            else "shm",
+                        )
                 continue
             _add_edge(
                 graph[task_idx][_DAGNodeOperationType.WRITE],

diff --git a/python/ray/dag/tests/experimental/test_accelerated_dag.py b/python/ray/dag/tests/experimental/test_accelerated_dag.py
index bda5ecfb41eb..a9f1b3955783 100644
--- a/python/ray/dag/tests/experimental/test_accelerated_dag.py
+++ b/python/ray/dag/tests/experimental/test_accelerated_dag.py
@@ -223,6 +223,7 @@ def test_inc_two_returns(ray_start_regular, single_fetch):
         dag = MultiOutputNode([o1, o2])

     compiled_dag = dag.experimental_compile()
+    compiled_dag.visualize(channel_details=True)
     for i in range(3):
         refs = compiled_dag.execute(1)
         if single_fetch:

[core][compiled graphs] Run fails only when using NCCL channels
### What happened + What you expected to happen

The following compiled graph works when I disable NCCL channel transfer: https://github.com/anmscale/verl/blob/ray_cg/verl/trainer/ppo/ray_trainer.py#L621-L662

However, when I enable NCCL transfer in these two steps I get an error:
https://github.com/anmscale/verl/blob/ray_cg/verl/trainer/ppo/ray_trainer.py#L638-L652

There are two different errors:
1) When I type-hint output of Actor-0 to Actor-0, which shouldn't be allowed at compile time but somehow works. I get the following error:
```
(main_task pid=2531538) Waiting for worker tasks to exit
(main_task pid=2531538) Teardown complete
Traceback (most recent call last):
  File "/home/ray/default/verl/verl/trainer/main_ppo.py", line 97, in main
    ray.get(main_task.remote(config))
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/auto_init_hook.py", line 21, in auto_init_wrapper
    return fn(*args, **kwargs)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/client_mode_hook.py", line 103, in wrapper
    return func(*args, **kwargs)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/worker.py", line 2763, in get
    values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/worker.py", line 914, in get_objects
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(AttributeError): ray::main_task() (pid=2531538, ip=10.0.0.27)
  File "/home/ray/default/verl/verl/trainer/main_ppo.py", line 195, in main_task
    trainer.fit()
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_20d7fa7e64594a715b9685aa68d7af73884e34d9/verl/trainer/ppo/ray_trainer.py", line 733, in fit
    metrics = self.run_compiled_graph(batch_dict)
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_20d7fa7e64594a715b9685aa68d7af73884e34d9/verl/trainer/ppo/ray_trainer.py", line 699, in run_compiled_graph
    dag_metrics = ray.get(dag_futures, timeout=100000)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 106, in get
    return _process_return_vals(return_vals, True)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 22, in _process_return_vals
    raise val.as_instanceof_cause()
ray.exceptions.RayTaskError(AttributeError): ray::WorkerDict.__ray_call__() (pid=2531917, ip=10.0.0.27)
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_20d7fa7e64594a715b9685aa68d7af73884e34d9/single_controller/ray/base.py", line 435, in func
    return getattr(self.worker_dict[key], name)(*args, **kwargs)
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_20d7fa7e64594a715b9685aa68d7af73884e34d9/single_controller/base/decorator.py", line 406, in inner
    return func(*args, **kwargs)
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_20d7fa7e64594a715b9685aa68d7af73884e34d9/verl/trainer/ppo/workers/fsdp_workers.py", line 926, in compute_scores_and_advantage
    data = DataProto.concat([data1, data2, data3, data4])
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_20d7fa7e64594a715b9685aa68d7af73884e34d9/verl/protocol.py", line 430, in concat
    batch_lst.append(batch.batch)
AttributeError: 'ray._raylet.MessagePackSerializedObject' object has no attribute 'batch'

```
2- When I enable nccl only across different Actors (i.e., disable it within Actor-0) I get the following:
```
(main_task pid=2794522) Tearing down compiled DAG
(main_task pid=2794522) Cancelling compiled worker on actor: Actor(create_colocated_worker_cls.<locals>.WorkerDict, 21547c8b679f8af158377dd190000000)
(main_task pid=2794522) Cancelling compiled worker on actor: Actor(create_colocated_worker_cls.<locals>.WorkerDict, 6106071252e14600d3031a3290000000)
(main_task pid=2794522) Cancelling compiled worker on actor: Actor(create_colocated_worker_cls.<locals>.WorkerDict, 0c16245c6e3308c33ac657f390000000)
(main_task pid=2794522) Cancelling compiled worker on actor: Actor(create_colocated_worker_cls.<locals>.WorkerDict, ffe689ca0de0981de9b764eb90000000)
(main_task pid=2794522) Exception ignored in: <function CompiledDAGRef.__del__ at 0x71bb6f199c10>
(main_task pid=2794522) Traceback (most recent call last):
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 93, in __del__
(main_task pid=2794522)     self.get()
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
(main_task pid=2794522)     return_vals = self._dag._execute_until(
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2122, in _execute_until
(main_task pid=2794522)     return self._get_execution_results(execution_index, channel_index)
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2047, in _get_execution_results
(main_task pid=2794522)     assert execution_index in self._result_buffer
(main_task pid=2794522) AssertionError:
(main_task pid=2794522) Exception ignored in: <function CompiledDAGRef.__del__ at 0x71bb6f199c10>
(main_task pid=2794522) Traceback (most recent call last):
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 93, in __del__
(main_task pid=2794522)     self.get()
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
(main_task pid=2794522)     return_vals = self._dag._execute_until(
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2122, in _execute_until
(main_task pid=2794522)     return self._get_execution_results(execution_index, channel_index)
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2047, in _get_execution_results
(main_task pid=2794522)     assert execution_index in self._result_buffer
(main_task pid=2794522) AssertionError:
(main_task pid=2794522) Exception ignored in: <function CompiledDAGRef.__del__ at 0x71bb6f199c10>
(main_task pid=2794522) Traceback (most recent call last):
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 93, in __del__
(main_task pid=2794522)     self.get()
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
(main_task pid=2794522)     return_vals = self._dag._execute_until(
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2122, in _execute_until
(main_task pid=2794522)     return self._get_execution_results(execution_index, channel_index)
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2047, in _get_execution_results
(main_task pid=2794522)     assert execution_index in self._result_buffer
(main_task pid=2794522) AssertionError:
(main_task pid=2794522) Exception ignored in: <function CompiledDAGRef.__del__ at 0x71bb6f199c10>
(main_task pid=2794522) Traceback (most recent call last):
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 93, in __del__
(main_task pid=2794522)     self.get()
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
(main_task pid=2794522)     return_vals = self._dag._execute_until(
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2122, in _execute_until
(main_task pid=2794522)     return self._get_execution_results(execution_index, channel_index)
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2047, in _get_execution_results
(main_task pid=2794522)     assert execution_index in self._result_buffer
(main_task pid=2794522) AssertionError:
(main_task pid=2794522) Exception ignored in: <function CompiledDAGRef.__del__ at 0x71bb6f199c10>
(main_task pid=2794522) Traceback (most recent call last):
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 93, in __del__
(main_task pid=2794522)     self.get()
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
(main_task pid=2794522)     return_vals = self._dag._execute_until(
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2122, in _execute_until
(main_task pid=2794522)     return self._get_execution_results(execution_index, channel_index)
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2047, in _get_execution_results
(main_task pid=2794522)     assert execution_index in self._result_buffer
(main_task pid=2794522) AssertionError:
(main_task pid=2794522) Exception ignored in: <function CompiledDAGRef.__del__ at 0x71bb6f199c10>
(main_task pid=2794522) Traceback (most recent call last):
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 93, in __del__
(main_task pid=2794522)     self.get()
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
(main_task pid=2794522)     return_vals = self._dag._execute_until(
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2122, in _execute_until
(main_task pid=2794522)     return self._get_execution_results(execution_index, channel_index)
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2047, in _get_execution_results
(main_task pid=2794522)     assert execution_index in self._result_buffer
(main_task pid=2794522) AssertionError:
(main_task pid=2794522) Exception ignored in: <function CompiledDAGRef.__del__ at 0x71bb6f199c10>
(main_task pid=2794522) Traceback (most recent call last):
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 93, in __del__
(main_task pid=2794522)     self.get()
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
(main_task pid=2794522)     return_vals = self._dag._execute_until(
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2122, in _execute_until
(main_task pid=2794522)     return self._get_execution_results(execution_index, channel_index)
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2047, in _get_execution_results
(main_task pid=2794522)     assert execution_index in self._result_buffer
(main_task pid=2794522) AssertionError:
(WorkerDict pid=2795038) /home/ray/anaconda3/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py:689: FutureWarning: FSDP.state_dict_type() and FSDP.set_state_dict_type() are being deprecated. Please use APIs, get_state_dict() and set_state_dict(), which can support different parallelisms, FSDP1, FSDP2, DDP. API doc: https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.state_dict.get_state_dict .Tutorial: https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html . [repeated 3x across cluster]
(WorkerDict pid=2795038)   warnings.warn( [repeated 3x across cluster]
(WorkerDict pid=2794846) TIMER: compute_values, CUDA time (s): 5.9334, perf_counter time (s): 5.9336
(WorkerDict pid=2794846) TIMER: union, CUDA time (s): 0.0001, perf_counter time (s): 0.0002
(WorkerDict pid=2794846) ERROR:2024-12-04 09:11:45,151:Compiled DAG task exited with exception
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 161, in do_exec_tasks
(WorkerDict pid=2794846)     done = tasks[operation.exec_task_idx].exec_operation(
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 644, in exec_operation
(WorkerDict pid=2794846)     return self._read(overlap_gpu_communication)
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 537, in _read
(WorkerDict pid=2794846)     input_data = self.input_reader.read()
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/common.py", line 296, in read
(WorkerDict pid=2794846)     outputs = self._read_list(timeout)
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/common.py", line 321, in _read_list
(WorkerDict pid=2794846)     results.append(c.read(timeout))
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/cached_channel.py", line 88, in read
(WorkerDict pid=2794846)     value = self._inner_channel.read(timeout)
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/shared_memory_channel.py", line 751, in read
(WorkerDict pid=2794846)     return self._channel_dict[actor_id].read(timeout)
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/intra_process_channel.py", line 64, in read
(WorkerDict pid=2794846)     return ctx.get_data(self._channel_id)
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/serialization_context.py", line 47, in get_data
(WorkerDict pid=2794846)     assert (
(WorkerDict pid=2794846) AssertionError: Channel 9d1c4b93-dd26-4f96-803d-2231c70dff4e does not exist in the buffer.
(WorkerDict pid=2795037) Destructing NCCL group on actor: Actor(create_colocated_worker_cls.<locals>.WorkerDict, ffe689ca0de0981de9b764eb90000000)
(main_task pid=2794522) Waiting for worker tasks to exit
(main_task pid=2794522) Teardown complete
Error executing job with overrides: ['data.train_files=/mnt/cluster_storage/data/gsm8k/train.parquet', 'data.val_files=/mnt/cluster_storage/data/gsm8k/test.parquet', 'data.train_batch_size=256', 'data.val_batch_size=512', 'data.max_prompt_length=512', 'data.max_response_length=512', 'actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat', 'actor_rollout_ref.actor.optim.lr=1e-6', 'actor_rollout_ref.actor.ppo_mini_batch_size=64', 'actor_rollout_ref.actor.ppo_micro_batch_size=8', 'actor_rollout_ref.actor.fsdp_config.param_offload=False', 'actor_rollout_ref.actor.fsdp_config.grad_offload=False', 'actor_rollout_ref.actor.fsdp_config.optimizer_offload=False', 'actor_rollout_ref.rollout.log_prob_micro_batch_size=32', 'actor_rollout_ref.rollout.tensor_model_parallel_size=4', 'actor_rollout_ref.rollout.name=vllm', 'actor_rollout_ref.rollout.gpu_memory_utilization=0.4', 'actor_rollout_ref.ref.log_prob_micro_batch_size=32', 'actor_rollout_ref.ref.fsdp_config.param_offload=True', 'critic.optim.lr=1e-5', 'critic.model.path=deepseek-ai/deepseek-llm-7b-chat', 'critic.model.enable_gradient_checkpointing=False', 'critic.ppo_micro_batch_size=8', 'critic.model.fsdp_config.param_offload=False', 'critic.model.fsdp_config.grad_offload=False', 'critic.model.fsdp_config.optimizer_offload=False', 'algorithm.kl_ctrl.kl_coef=0.001', 'trainer.critic_warmup=0', 'trainer.logger=[console]', 'trainer.project_name=verl_example_gsm8k', 'trainer.experiment_name=deepseek_llm_7b_function_rm', 'trainer.n_gpus_per_node=4', 'trainer.nnodes=1', 'trainer.save_freq=-1', 'trainer.test_freq=10', 'trainer.total_epochs=5', 'trainer.default_hdfs_dir=/mnt/local_storage/experiments/gsm8k/ppo/deepseek_llm_7b_function_rm', 'trainer.use_rcg=True']
Traceback (most recent call last):
  File "/home/ray/default/verl/verl/trainer/main_ppo.py", line 97, in main
    ray.get(main_task.remote(config))
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/auto_init_hook.py", line 21, in auto_init_wrapper
    return fn(*args, **kwargs)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/client_mode_hook.py", line 103, in wrapper
    return func(*args, **kwargs)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/worker.py", line 2763, in get
    values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/worker.py", line 914, in get_objects
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(RayChannelError): ray::main_task() (pid=2794522, ip=10.0.0.27)
  File "/home/ray/default/verl/verl/trainer/main_ppo.py", line 195, in main_task
    trainer.fit()
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_326f52a2699eb2dc53e27c8a8ba5a9863db795ec/verl/trainer/ppo/ray_trainer.py", line 733, in fit
    metrics = self.run_compiled_graph(batch_dict)
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_326f52a2699eb2dc53e27c8a8ba5a9863db795ec/verl/trainer/ppo/ray_trainer.py", line 699, in run_compiled_graph
    dag_metrics = ray.get(dag_futures, timeout=100000)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
    return_vals = self._dag._execute_until(
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2115, in _execute_until
    self._dag_output_fetcher.read(timeout),
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/common.py", line 296, in read
    outputs = self._read_list(timeout)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/common.py", line 321, in _read_list
    results.append(c.read(timeout))
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/shared_memory_channel.py", line 751, in read
    return self._channel_dict[actor_id].read(timeout)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/shared_memory_channel.py", line 602, in read
    output = self._buffers[self._next_read_index].read(timeout)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/shared_memory_channel.py", line 485, in read
    ret = self._worker.get_objects(
  File "python/ray/includes/common.pxi", line 100, in ray._raylet.check_status
ray.exceptions.RayChannelError: System error: Channel closed.
```

### Versions / Dependencies

```
>>> ray.__version__
'3.0.0.dev0'
>>> ray.__commit__
'62d59ba60823ceb7b2fb9a121c40ff58042be166'
```

### Reproduction script

https://github.com/anmscale/verl/pull/2

### Issue Severity

High: It blocks me from completing my task.

I was able to reproduce the bug in the following standalone script. Note that it works fine when `nccl=False`.
https://github.com/anmscale/verl/blob/ray_cg/examples/ray/rcg_bug.py

For context, here's the workload I am trying to achieve:
![image](https://github.com/user-attachments/assets/98415937-0d13-4cc6-b517-157e15809cfc)

I think the issue is this part:

```
def scatter_from_rank_zero(
    workers: List[Worker], data: DataProto, nccl: bool = False
) -> Tuple[DataProto]:
    n_shards = len(workers)
    sharded_data = workers[0].chunk.options(num_returns=n_shards).bind(data, n_shards)
    type_hinted_data = [sharded_data[0]]
    for data in sharded_data[1:]:
        type_hinted_data.append(
            data.with_type_hint(TorchTensorType(transport="nccl")) if nccl else data
        )
    return tuple(type_hinted_data)
```

The type hint put on partial returns (e.g., when num_returns=n) is not currently respected. I will look into a fix.
Debugged further, the issue is missing data dependency when `num_returns` is used.

```(Worker pid=2902899) AssertionError: Channel cc21f308-88ff-4869-ac22-30130fc235aa does not exist in the buffer. ```
![acompiled_graph](https://github.com/user-attachments/assets/5bbc7146-81c6-4f7b-be93-08a432ecd586)

![acompiled_graph_schedule](https://github.com/user-attachments/assets/288b287f-9695-430f-8d7c-9ee5742e3350)


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

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,
         )

[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

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.