diff --git a/python/ray/_private/log.py b/python/ray/_private/log.py
index 73ce2d3ef329..69242c722233 100644
--- a/python/ray/_private/log.py
+++ b/python/ray/_private/log.py
@@ -1,5 +1,4 @@
 import logging
-from logging.config import dictConfig
 import threading
 from typing import Union

@@ -74,44 +73,19 @@ def generate_logging_config():
             return
         logger_initialized = True

-        formatters = {
-            "plain": {
-                "format": (
-                    "%(asctime)s\t%(levelname)s %(filename)s:%(lineno)s -- %(message)s"
-                ),
-            },
-        }
-
-        handlers = {
-            "default": {
-                "()": PlainRayHandler,
-                "formatter": "plain",
-            }
-        }
-
-        loggers = {
-            # Default ray logger; any log message that gets propagated here will be
-            # logged to the console. Disable propagation, as many users will use
-            # basicConfig to set up a default handler. If so, logs will be
-            # printed twice unless we prevent propagation here.
-            "ray": {
-                "level": "INFO",
-                "handlers": ["default"],
-                "propagate": False,
-            },
-            # Special handling for ray.rllib: only warning-level messages passed through
-            # See https://github.com/ray-project/ray/pull/31858 for related PR
-            "ray.rllib": {
-                "level": "WARN",
-            },
-        }
-
-        dictConfig(
-            {
-                "version": 1,
-                "formatters": formatters,
-                "handlers": handlers,
-                "loggers": loggers,
-                "disable_existing_loggers": False,
-            }
+        plain_formatter = logging.Formatter(
+            "%(asctime)s\t%(levelname)s %(filename)s:%(lineno)s -- %(message)s"
         )
+
+        default_handler = PlainRayHandler()
+        default_handler.setFormatter(plain_formatter)
+
+        ray_logger = logging.getLogger("ray")
+        ray_logger.setLevel(logging.INFO)
+        ray_logger.addHandler(default_handler)
+        ray_logger.propagate = False
+
+        # Special handling for ray.rllib: only warning-level messages passed through
+        # See https://github.com/ray-project/ray/pull/31858 for related PR
+        rllib_logger = logging.getLogger("ray.rllib")
+        rllib_logger.setLevel(logging.WARN)
diff --git a/python/ray/_private/ray_logging/__init__.py b/python/ray/_private/ray_logging/__init__.py
index 5699251e31ce..768bc716c7c3 100644
--- a/python/ray/_private/ray_logging/__init__.py
+++ b/python/ray/_private/ray_logging/__init__.py
@@ -1,6 +1,7 @@
 import colorama
 from dataclasses import dataclass
 import logging
+import logging.handlers
 import os
 import re
 import sys
diff --git a/python/ray/_private/ray_logging/default_impl.py b/python/ray/_private/ray_logging/default_impl.py
index 2d50673c2851..3901baaa956f 100644
--- a/python/ray/_private/ray_logging/default_impl.py
+++ b/python/ray/_private/ray_logging/default_impl.py
@@ -1,4 +1,4 @@
-def get_dict_config_provider():
-    from ray._private.ray_logging.logging_config import DefaultDictConfigProvider
+def get_logging_configurator():
+    from ray._private.ray_logging.logging_config import DefaultLoggingConfigurator

-    return DefaultDictConfigProvider()
+    return DefaultLoggingConfigurator()
diff --git a/python/ray/_private/ray_logging/logging_config.py b/python/ray/_private/ray_logging/logging_config.py
index 6788571f9e6a..87cf819abbcf 100644
--- a/python/ray/_private/ray_logging/logging_config.py
+++ b/python/ray/_private/ray_logging/logging_config.py
@@ -12,68 +12,47 @@
 import time

-class DictConfigProvider(ABC):
+class LoggingConfigurator(ABC):
     @abstractmethod
     def get_supported_encodings(self) -> Set[str]:
         raise NotImplementedError

     @abstractmethod
-    def get_dict_config(self, encoding: str, log_level: str) -> dict:
+    def configure_logging(self, encoding: str, log_level: str):
         raise NotImplementedError

-class DefaultDictConfigProvider(DictConfigProvider):
+class DefaultLoggingConfigurator(LoggingConfigurator):
     def __init__(self):
-        self._dict_configs = {
-            "TEXT": lambda log_level: {
-                "version": 1,
-                "disable_existing_loggers": False,
-                "formatters": {
-                    "text": {
-                        "()": (
-                            f"{TextFormatter.__module__}."
-                            f"{TextFormatter.__qualname__}"
-                        ),
-                    },
-                },
-                "filters": {
-                    "core_context": {
-                        "()": (
-                            f"{CoreContextFilter.__module__}."
-                            f"{CoreContextFilter.__qualname__}"
-                        ),
-                    },
-                },
-                "handlers": {
-                    "console": {
-                        "level": log_level,
-                        "class": "logging.StreamHandler",
-                        "formatter": "text",
-                        "filters": ["core_context"],
-                    },
-                },
-                "root": {
-                    "level": log_level,
-                    "handlers": ["console"],
-                },
-                "loggers": {
-                    "ray": {
-                        "level": log_level,
-                        "handlers": ["console"],
-                        "propagate": False,
-                    }
-                },
-            }
+        self._encoding_to_formatter = {
+            "TEXT": TextFormatter(),
         }

     def get_supported_encodings(self) -> Set[str]:
-        return self._dict_configs.keys()
+        return self._encoding_to_formatter.keys()

-    def get_dict_config(self, encoding: str, log_level: str) -> dict:
-        return self._dict_configs[encoding](log_level)
+    def configure_logging(self, encoding: str, log_level: str):
+        formatter = self._encoding_to_formatter[encoding]
+        core_context_filter = CoreContextFilter()
+        handler = logging.StreamHandler()
+        handler.setLevel(log_level)
+        handler.setFormatter(formatter)
+        handler.addFilter(core_context_filter)

+        root_logger = logging.getLogger()
+        root_logger.setLevel(log_level)
+        root_logger.addHandler(handler)

-_dict_config_provider: DictConfigProvider = default_impl.get_dict_config_provider()
+        ray_logger = logging.getLogger("ray")
+        ray_logger.setLevel(log_level)
+        # Remove all existing handlers added by `ray/__init__.py`.
+        for h in ray_logger.handlers[:]:
+            ray_logger.removeHandler(h)
+        ray_logger.addHandler(handler)
+        ray_logger.propagate = False
+
+
+_logging_configurator: LoggingConfigurator = default_impl.get_logging_configurator()

 @PublicAPI(stability="alpha")
@@ -84,19 +63,16 @@ class LoggingConfig:
     log_level: str = "INFO"

     def __post_init__(self):
-        if self.encoding not in _dict_config_provider.get_supported_encodings():
+        if self.encoding not in _logging_configurator.get_supported_encodings():
             raise ValueError(
                 f"Invalid encoding type: {self.encoding}. "
                 "Valid encoding types are: "
-                f"{list(_dict_config_provider.get_supported_encodings())}"
+                f"{list(_logging_configurator.get_supported_encodings())}"
             )

-    def _get_dict_config(self) -> dict:
-        """Get the logging configuration based on the encoding type.
-        Returns:
-            dict: The logging configuration.
-        """
-        return _dict_config_provider.get_dict_config(self.encoding, self.log_level)
+    def _configure_logging(self):
+        """Set up the logging configuration for the current process."""
+        _logging_configurator.configure_logging(self.encoding, self.log_level)

     def _setup_log_record_factory(self):
         old_factory = logging.getLogRecordFactory()
@@ -121,7 +97,7 @@ def record_factory(*args, **kwargs):
     def _apply(self):
         """Set up both the LogRecord factory and the logging configuration."""
         self._setup_log_record_factory()
-        logging.config.dictConfig(self._get_dict_config())
+        self._configure_logging()

 LoggingConfig.__doc__ = f"""
@@ -153,7 +129,7 @@ def f():

     Args:
         encoding: Encoding type for the logs. The valid values are
-            {list(_dict_config_provider.get_supported_encodings())}
+            {list(_logging_configurator.get_supported_encodings())}
         log_level: Log level for the logs. Defaults to 'INFO'. You can set
             it to 'DEBUG' to receive more detailed debug logs.
     """  # noqa: E501

diff --git a/python/ray/tests/test_logging_2.py b/python/ray/tests/test_logging_2.py
index 9d5be165f9ac..988ed482e2bc 100644
--- a/python/ray/tests/test_logging_2.py
+++ b/python/ray/tests/test_logging_2.py
@@ -207,17 +207,9 @@ def test_record_with_exception(self):
             assert s in formatted

-class TestLoggingConfig:
-    def test_log_level(self):
-        log_level = "DEBUG"
-        logging_config = LoggingConfig(log_level=log_level)
-        dict_config = logging_config._get_dict_config()
-        assert dict_config["handlers"]["console"]["level"] == log_level
-        assert dict_config["root"]["level"] == log_level
-
-    def test_invalid_dict_config(self):
-        with pytest.raises(ValueError):
-            LoggingConfig(encoding="INVALID")._get_dict_config()
+def test_invalid_encoding():
+    with pytest.raises(ValueError):
+        LoggingConfig(encoding="INVALID")

 class TestTextModeE2E:
@@ -477,6 +469,27 @@ def print_message(self):
     assert "(MyActor pid=" not in stderr

+def test_configure_both_structured_logging_and_lib_logging(shutdown_only):
+    """
+    Configure the `ray.test` logger. Then, configure the `root` and `ray`
+    loggers in `ray.init()`. Ensure that the `ray.test` logger is not affected.
+    """
+    script = """
+import ray
+import logging
+
+old_test_logger = logging.getLogger("ray.test")
+assert old_test_logger.getEffectiveLevel() != logging.DEBUG
+old_test_logger.setLevel(logging.DEBUG)
+
+ray.init(logging_config=ray.LoggingConfig(encoding="TEXT", log_level="INFO"))
+
+new_test_logger = logging.getLogger("ray.test")
+assert old_test_logger.getEffectiveLevel() == logging.DEBUG
+"""
+    run_string_as_driver(script)
+
+
 if __name__ == "__main__":
     if os.environ.get("PARALLEL_CI"):
         sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))

[Core] Ray Core / Ray Data logging configuration leads to unexpected behavior
### What happened + What you expected to happen

We use `logging.config.dictConfig(config)` to configure the ray data logger  ([here](https://github.com/ray-project/ray/blob/4f6a419467885c295700f2db6375a298bc575fc2/python/ray/data/_internal/logging.py#L176)), but this is also how ray core configures the ray logger ([here](https://github.com/ray-project/ray/blob/4f6a419467885c295700f2db6375a298bc575fc2/python/ray/_private/ray_logging/logging_config.py#L124)).

For both of these logging configs, we use `disable_existing_loggers: False`. The behavior for this is described as ([logging docs](https://docs.python.org/3/library/logging.config.html#dictionary-schema-details)):
> disable_existing_loggers - whether any existing non-root loggers are to be disabled. This setting mirrors the parameter of the same name in [fileConfig()](https://docs.python.org/3/library/logging.config.html#logging.config.fileConfig). If absent, this parameter defaults to True. This value is ignored if incremental is True.

This description makes it seem like these two logging calls are commutative (regardless of ordering they will produce the same result), but that is not exactly how the python logging module works. If we configure the ray module logger then the ray.data module logger, the results are expected and both are configured. If we instead configure the ray.data module then configure the ray module logger, then the ray.data logging configuration is clobbered. This happens because when configuring the parent logger of a module (e.g. ray module logger is the parent logger of the ray.data module logger), the various handlers associated with the child logger are not guaranteed to be preserved.

Our end goal should be a state where the call order of the logging configurations should not affect the logging behavior.

### Versions / Dependencies

ray==2.39.0

### Reproduction script

```python
import ray
import logging

def report_logger(logger):
    # Collect this logger and its parents
    loggers = []
    current_logger = logger
    while current_logger:
        loggers.append(current_logger)
        if not current_logger.parent or current_logger.parent == current_logger:
            break
        current_logger = current_logger.parent

    # Report the configuration of each logger in the hierarchy
    print(f"Logging configuration for '{logger.name}' and its hierarchy:")
    for log in reversed(loggers):  # Start from the root and go down to the given logger
        print(f"\nLogger: {log.name or 'root'} (Level: {logging.getLevelName(log.level)})")
        if log.handlers:
            print("  Handlers:")
            for handler in log.handlers:
                print(f"    - {handler.__class__.__name__} (Level: {logging.getLevelName(handler.level)})")
        else:
            print("  No handlers configured")

print("BEFORE")
report_logger(logging.getLogger("ray.data"))
print()

import ray.data
ray.init(logging_config=ray.LoggingConfig(encoding="JSON", log_level="INFO"))

print("AFTER:")
report_logger(logging.getLogger("ray.data"))
```

### Issue Severity

High: It blocks me from completing my task.

To be clear about the expected behavior some examples / modifications of the repro:

### 1. Ray Data Import then Ray Core Init [UNEXPECTED BEHAVIOR]
Script:
```python
print("BEFORE")
report_logger(logging.getLogger("ray.data"))
print()

import ray.data
ray.init(logging_config=ray.LoggingConfig(encoding="JSON", log_level="INFO"))

print("AFTER:")
report_logger(logging.getLogger("ray.data"))
```
Output:
```
BEFORE
Logging configuration for 'ray.data' and its hierarchy:

Logger: root (Level: WARNING)
  No handlers configured

Logger: ray (Level: INFO)
  Handlers:
    - PlainRayHandler (Level: NOTSET)

Logger: ray.data (Level: NOTSET)
  No handlers configured

AFTER:
Logging configuration for 'ray.data' and its hierarchy:

Logger: root (Level: INFO)
  Handlers:
    - StreamHandler (Level: INFO)

Logger: ray (Level: INFO)
  Handlers:
    - StreamHandler (Level: INFO)

Logger: ray.data (Level: NOTSET)
  No handlers configured

```
Notes:
Ray Data configuration ignored, only Ray Core initialization respected.

### 2. Ray Core Init then Ray Data Import [EXPECTED BEHAVIOR]
Script:
```python
print("BEFORE")
report_logger(logging.getLogger("ray.data"))
print()

ray.init(logging_config=ray.LoggingConfig(encoding="JSON", log_level="INFO"))
import ray.data

print("AFTER:")
report_logger(logging.getLogger("ray.data"))
```
Output:
```
BEFORE
Logging configuration for 'ray.data' and its hierarchy:

Logger: root (Level: WARNING)
  No handlers configured

Logger: ray (Level: INFO)
  Handlers:
    - PlainRayHandler (Level: NOTSET)

Logger: ray.data (Level: NOTSET)
  No handlers configured

AFTER:
Logging configuration for 'ray.data' and its hierarchy:

Logger: root (Level: INFO)
  Handlers:
    - StreamHandler (Level: INFO)

Logger: ray (Level: INFO)
  Handlers:
    - StreamHandler (Level: INFO)

Logger: ray.data (Level: DEBUG)
  Handlers:
    - SessionFileHandler (Level: NOTSET)
    - PlainRayHandler (Level: INFO)
    - SessionFileHandler (Level: ERROR)
```
Notes:
Ray Core configuration is done first, Ray Data configuration is done second, both are respected.

### 3. Only Ray Data Import [EXPECTED BEHAVIOR]
Script:
```python
print("BEFORE")
report_logger(logging.getLogger("ray.data"))
print()

import ray.data

print("AFTER:")
report_logger(logging.getLogger("ray.data"))
```
Output:
```
BEFORE
Logging configuration for 'ray.data' and its hierarchy:

Logger: root (Level: WARNING)
  No handlers configured

Logger: ray (Level: INFO)
  Handlers:
    - PlainRayHandler (Level: NOTSET)

Logger: ray.data (Level: NOTSET)
  No handlers configured

AFTER:
Logging configuration for 'ray.data' and its hierarchy:

Logger: root (Level: WARNING)
  No handlers configured

Logger: ray (Level: INFO)
  Handlers:
    - PlainRayHandler (Level: NOTSET)

Logger: ray.data (Level: DEBUG)
  Handlers:
    - SessionFileHandler (Level: NOTSET)
    - PlainRayHandler (Level: INFO)
    - SessionFileHandler (Level: ERROR)
```
Notes:
Correctly configures the Ray Data logger.

### 4. Only Ray Core Init [EXPECTED BEHAVIOR]
Script:
```python
print("BEFORE")
report_logger(logging.getLogger("ray.data"))
print()

ray.init(logging_config=ray.LoggingConfig(encoding="JSON", log_level="INFO"))

print("AFTER:")
report_logger(logging.getLogger("ray.data"))
```
Output:
```
BEFORE
Logging configuration for 'ray.data' and its hierarchy:

Logger: root (Level: WARNING)
  No handlers configured

Logger: ray (Level: INFO)
  Handlers:
    - PlainRayHandler (Level: NOTSET)

Logger: ray.data (Level: NOTSET)
  No handlers configured

AFTER:
Logging configuration for 'ray.data' and its hierarchy:

Logger: root (Level: INFO)
  Handlers:
    - StreamHandler (Level: INFO)

Logger: ray (Level: INFO)
  Handlers:
    - StreamHandler (Level: INFO)

Logger: ray.data (Level: NOTSET)
  No handlers configured

```
Notes:
Correctly configures the Ray Core logger.

@omatthew98 why don't we just do

1. Take the current config
2. Merge in provided one
3. Set the config back
Update from some more offline discussion with @alexeykudinkin:
For now, ensuring that the import / happens after the initialization (e.g. ensuring that the ray module logger is configured before the ray.data module logger), should be an immediate workaround. We should still try to improve the behavior as this is a footgun that will result in unexpected results and missing logs that will hinder debugging efforts for users. Ideally we would do what was suggested above, but unfortunately python's logging module does not provide a native way to retrieve the existing logging configuration as a dictionary, so we would have to manually inspect and reconstruct the dict before merging the provided config which would be somewhat hacky.

An alternative to this would be to just wrap this configuration at the ray core level by:
1. Persist configuration as root level field of the core logging module
2. Create util `configureLogging` and have it be called from both core and other libraries
3. Within the `configureLogging` function, merge the provided dictionary with the existing dictionary configuration and then update the logging config with the merged config.
I am trying to understand the structured logging behavior as well. One of my question is
#1. Ray Data Import then Ray Core Init [UNEXPECTED BEHAVIOR] will wipe out the `ray.data` logger configuration; but why
#2. Ray Core Init then Ray Data Import [EXPECTED BEHAVIOR] maintains the expected logger configurations of both logger?
It seems by going in order #2, we got the correct logger config for all three loggers: i.e., `root` by ray.init, `ray` by ray.init, and finally `ray.data` by import ray data
> I am trying to understand the structured logging behavior as well. One of my question is #1. Ray Data Import then Ray Core Init [UNEXPECTED BEHAVIOR] will wipe out the `ray.data` logger configuration; but why #2. Ray Core Init then Ray Data Import [EXPECTED BEHAVIOR] maintains the expected logger configurations of both logger? It seems by going in order #2, we got the correct logger config for all three loggers: i.e., `root` by ray.init, `ray` by ray.init, and finally `ray.data` by import ray data

Yeah I think that understanding is correct. From what I have read, configuring the parent logger after the child logger (so the order of #1), the result will be only the parent logger is configured. If you configure the child logger after the parent logger (so the order of #2), then the two loggers will be configured as expected. I think this is is just an implementation detail for python's logging module that doesn't seem particularly well documented.


diff --git a/python/ray/experimental/channel/serialization_context.py b/python/ray/experimental/channel/serialization_context.py
index 613f923b7a6c..4aa5f280b51d 100644
--- a/python/ray/experimental/channel/serialization_context.py
+++ b/python/ray/experimental/channel/serialization_context.py
@@ -97,7 +97,11 @@ def serialize_tensor(self, tensor: "torch.Tensor") -> Union[int, "np.ndarray"]:

         return self.serialize_to_numpy(tensor)

-    def serialize_to_numpy(self, tensor: "torch.Tensor") -> "np.ndarray":
+    def serialize_to_numpy(
+        self, tensor: "torch.Tensor"
+    ) -> Tuple["np.ndarray", "torch.dtype"]:
+        import torch
+
         # Transfer through Ray's shared memory store for now.
         # TODO(swang): This requires two copies, one to transfer from GPU to
         # CPU and another from CPU to shared memory. Ideally we should elide
@@ -106,7 +110,10 @@ def serialize_to_numpy(self, tensor: "torch.Tensor") -> "np.ndarray":
         if tensor.device.type == "cuda":
             tensor = tensor.to("cpu")

-        return tensor.numpy()
+        # Numpy does not have an equivalent dtype for all torch dtypes, so
+        # instead of casting directly to numpy, we first use a view with a
+        # common dtype and then view as numpy array.
+        return (tensor.view(torch.uint8).numpy(), tensor.dtype)

     def deserialize_tensor(self, val: Union["np.ndarray", int]):
         # Found a placeholder for a tensor that was serialized via NCCL.
@@ -119,13 +126,17 @@ def deserialize_tensor(self, val: Union["np.ndarray", int]):

         return self.deserialize_from_numpy(val)

-    def deserialize_from_numpy(self, np_array: "np.ndarray"):
+    def deserialize_from_numpy(
+        self, np_array_dtype: Tuple["np.ndarray", "torch.dtype"]
+    ):
         import torch

         from ray.experimental.channel import ChannelContext

         ctx = ChannelContext.get_current()

+        np_array, dtype = np_array_dtype
+
         # TODO(swang): Support local P2P transfers if available.
         # If there is a GPU assigned to this worker, move it there.
         if ctx.torch_device is not None and ctx.torch_device.type == "cuda":
@@ -133,7 +144,7 @@ def deserialize_from_numpy(self, np_array: "np.ndarray"):
             def convert_numpy_to_tensor(np_array, ctx):
                 # It does zero-copy convert np_array inside shared memroy to
                 # a tensor. Since we move data to GPU immediately, it is safe.
-                cpu_tensor = torch.from_numpy(np_array)
+                cpu_tensor = torch.from_numpy(np_array).view(dtype)
                 return cpu_tensor.to(device=ctx.torch_device)

             global _TORCH_WARNING_FILTER_ACTIVATE
@@ -160,4 +171,4 @@ def convert_numpy_to_tensor(np_array, ctx):
         # TODO(swang): Use zero-copy from_numpy() if np_array.flags.writeable
         # is True. This is safe to set when deserializing np_array if the
         # upstream task has num_readers=1.
-        return torch.tensor(np_array, device=ctx.torch_device)
+        return torch.tensor(np_array, device=ctx.torch_device).view(dtype)

diff --git a/python/ray/tests/test_channel.py b/python/ray/tests/test_channel.py
index 69127ee157a8..95fc659c5a32 100644
--- a/python/ray/tests/test_channel.py
+++ b/python/ray/tests/test_channel.py
@@ -8,11 +8,13 @@

 import numpy as np
 import pytest
+import torch

 import ray
 import ray.cluster_utils
 import ray.exceptions
 import ray.experimental.channel as ray_channel
+from ray.experimental.channel.torch_tensor_type import TorchTensorType
 from ray.exceptions import RayChannelError, RayChannelTimeoutError
 from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
 from ray.dag.compiled_dag_node import CompiledDAG
@@ -1407,6 +1409,20 @@ def write(self, i, timeout=None) -> bool:
         )

+def test_torch_dtype():
+    typ = TorchTensorType()
+    typ.register_custom_serializer()
+
+    t = torch.randn(5, 5, dtype=torch.bfloat16)
+    with pytest.raises(TypeError):
+        t.numpy()
+
+    ref = ray.put(t)
+    t_out = ray.get(ref)
+    assert (t == t_out).all()
+    assert t_out.dtype == t.dtype
+
+
 if __name__ == "__main__":
     if os.environ.get("PARALLEL_CI"):
         sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))

[core][compiled graphs] Support all torch.dtypes for tensors sent through shared memory channels
### What happened + What you expected to happen

Got the following error when experimenting with OpenRLHF:

**TypeError: Got unsupported ScalarType BFloat16** :

```
(CriticModelRayActor pid=33730)   cpu_tensor = torch.from_numpy(np_array) [repeated 2x across cluster]
(RewardModelRayActor pid=33731) Traceback (most recent call last): [repeated 4x across cluster]
(RewardModelRayActor pid=33731) 2024-10-21 09:10:49 ERROR    Compiled DAG task exited with exception
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/experimental/channel/shared_memory_channel.py", line 469, in write [repeated 5x across cluster]
(RewardModelRayActor pid=33731)     serialized_value = self._worker.get_serialization_context().serialize(
(RewardModelRayActor pid=33731)                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/experimental/channel/torch_tensor_type.py", line 108, in serialize [repeated 2x across cluster]
(RewardModelRayActor pid=33731)     return self._serialize_to_msgpack(value)
(RewardModelRayActor pid=33731)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/_private/serialization.py", line 497, in _serialize_to_msgpack
(RewardModelRayActor pid=33731)     pickle5_serialized_object = self._serialize_to_pickle5(
(RewardModelRayActor pid=33731)                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/_private/serialization.py", line 439, in _serialize_to_pickle5 [repeated 2x across cluster]
(RewardModelRayActor pid=33731)     raise e
(RewardModelRayActor pid=33731)     inband = pickle.dumps(
(RewardModelRayActor pid=33731)            ^^^^^^^^^^^^^ [repeated 2x across cluster]
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/cloudpickle/cloudpickle.py", line 1479, in dumps
(RewardModelRayActor pid=33731)     cp.dump(obj)
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/cloudpickle/cloudpickle.py", line 1245, in dump
(RewardModelRayActor pid=33731)     return super().dump(obj)
(RewardModelRayActor pid=33731)            ^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/_private/serialization.py", line 175, in _CloudPicklerReducer
(RewardModelRayActor pid=33731)     return custom_deserializer, (custom_serializer(obj),)
(RewardModelRayActor pid=33731)                                  ^^^^^^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)     return ctx.serialization_context.serialize_tensor(t)
(RewardModelRayActor pid=33731)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/experimental/channel/serialization_context.py", line 77, in serialize_tensor
(RewardModelRayActor pid=33731)     return self.serialize_to_numpy(tensor)
(RewardModelRayActor pid=33731)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/experimental/channel/serialization_context.py", line 88, in serialize_to_numpy
(RewardModelRayActor pid=33731)     return tensor.numpy()
(RewardModelRayActor pid=33731)            ^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731) TypeError: Got unsupported ScalarType BFloat16
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/dag/compiled_dag_node.py", line 118, in do_exec_tasks
(RewardModelRayActor pid=33731)     done = tasks[operation.exec_task_idx].exec_operation(
(RewardModelRayActor pid=33731)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/dag/compiled_dag_node.py", line 547, in exec_operation
(RewardModelRayActor pid=33731)     return self._write()
(RewardModelRayActor pid=33731)   File "/home/ray/anaconda3/lib/python3.11/site-packages/ray/dag/compiled_dag_node.py", line 517, in _write
(RewardModelRayActor pid=33731)     self.output_writer.write(output_val)
(RewardModelRayActor pid=33731)     channel.write(val_i, timeout)
(RewardModelRayActor pid=33731)     channel.write(value, timeout)
(RewardModelRayActor pid=33731)     self._buffers[self._next_write_index].write(value, timeout)
```

### Versions / Dependencies

head

### Reproduction script

will add later

### Issue Severity

None

Another data point is that Pytorch has a number of other data types that aren't natively supported in numpy (like float8, etc). I'm not particularly sure of the details for which `tensor` was converted to a numpy array here but it would be best to not assume any interoperability with numpy and keep tensors as tensors during serialization.
Hmm really what we need is to be able to zero-copy deserialize torch.Tensors. Using numpy was easiest at the time but this is indeed an issue. Maybe [arrow](https://arrow.apache.org/docs/cpp/api/tensor.html#tensors) is an option?
Hmm actually seems like this could be supported by using torch.Tensor.view with a uint8 dtype:

```python
In [99]: t
Out[99]:
tensor([ 1.8750, -0.6875, -1.1250, -1.3750,  1.3750, -1.1250,  0.4688, -0.4062,
         0.8750, -1.7500], dtype=torch.float8_e4m3fn)

In [104]: torch.as_tensor(t.view(torch.uint8).numpy()).view(torch.float8_e4m3fn)
Out[104]:
tensor([ 1.8750, -0.6875, -1.1250, -1.3750,  1.3750, -1.1250,  0.4688, -0.4062,
         0.8750, -1.7500], dtype=torch.float8_e4m3fn)

```
Hi @stephanie-wang,

I am interested in taking up this issue.
> Hi @stephanie-wang,
>
> I am interested in taking up this issue.

Thanks, feel free to open a PR. Will assign it to you for now and revisit in a week or so.
Thanks, do you have suggestions on reproducing the issue on my local?
You want a DAG that looks something like:
```python
t = A.return_tensor(...).with_type_hint(TorchTensorType())
dag = B.read_tensor(t)
```
And A should return a tensor with a dtype unsupported by numpy.
I am new to the eco-system. I'm sorry if this is an obvious question. What do `A` and `B` represent in your previous comment? I tried creating an example using the "Getting Started" docs. However, I am struggling to reproduce the issue.

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

@ray.remote
def return_tensor():
    return torch.Tensor([1.0, 2.0])

@ray.remote
def read_tensor(a):
    return a

dag = read_tensor.bind(return_tensor.bind().with_type_hint(TorchTensorType()))

print(workflow.run(dag))
```

Am I on the right track? Also, if I try to specify the `dtype` param in the `torch.Tensor` init, the program fails with a weird torch data type error. I am not sure if that's related.
Hi @sahilgupta2105, yes we don't have great docs for compiled graphs yet. Please work through [the developer guide](https://docs.google.com/document/d/1xDs-yIBFLXuqUDLou-_KELERkzWrmGWri6p_xG18SH4/edit?tab=t.0#heading=h.kzrrbhon06yz) first and comment here if you still have questions.
Thanks, the developer guide is helpful. I was able to reproduce the issue.

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

@ray.remote
class Actor:
  def process(self, tensor: torch.Tensor):
    return tensor.shape

actor = Actor.remote()

with ray.dag.InputNode() as inp:
  inp = inp.with_type_hint(TorchTensorType())
  dag = actor.process.bind(inp)

  dag = dag.experimental_compile()
  print(ray.get(dag.execute(torch.zeros(10, dtype=torch.float8_e4m3fn))))
```
> Hmm actually seems like this could be supported by using torch.Tensor.view with a uint8 dtype:
>
> ```python
> In [99]: t
> Out[99]:
> tensor([ 1.8750, -0.6875, -1.1250, -1.3750,  1.3750, -1.1250,  0.4688, -0.4062,
>          0.8750, -1.7500], dtype=torch.float8_e4m3fn)
>
> In [104]: torch.as_tensor(t.view(torch.uint8).numpy()).view(torch.float8_e4m3fn)
> Out[104]:
> tensor([ 1.8750, -0.6875, -1.1250, -1.3750,  1.3750, -1.1250,  0.4688, -0.4062,
>          0.8750, -1.7500], dtype=torch.float8_e4m3fn)
> ```

@stephanie-wang to be sure, you meant type-casting the torch tensor to the unsigned int type with the appropriate precision, right? E.g. float8 -> uint8, float16 -> uint16, ...
No, cast everything to uint8 (because it's a supported np dtype and the lowest possible precision) and then cast back on the receiving side.
Gotcha. For my understanding, wouldn't casting everything to uint8 cause a loss of information if the input tensor is of higher precision?
@stephanie-wang It seems like the casting operations only update the metadata, which ensures that information is not lost when `floatXX` is converted to `uint8`. However, to ensure data integrity, we need information about the `dtype` of the original array. Numpy arrays support [custom metadata](https://numpy.org/doc/stable/reference/generated/numpy.dtype.metadata.html). Shall we use that to store the `dtype` on serialization for deserialization to ensure the correct format?
It would be better to not rely on a third-party API for passing the dtype. You can pass it through Ray instead. Check out how [TorchTensorType](https://github.com/ray-project/ray/blob/master/python/ray/experimental/channel/torch_tensor_type.py) is used.
I dug more into the code. Can you confirm my understanding before I send out a PR?

- `TorchTensorType` uses a serializer from the serialization context that internally calls the `.numpy()`.
- Each `TorchTensorType` hint is unique per DAG node and contains the type in `_dtype` attribute.
- We can declare a custom serializer (skip calling the serializer from the serialization context) and use the `_dtype` attribute to maintain data integrity.

If you agree so far, do I need special handling for the "AUTO" dtype?
I don't think you need to touch TorchTensorType at all. Should be enough to just modify the existing custom serialization function.


diff --git a/python/ray/data/grouped_data.py b/python/ray/data/grouped_data.py
index 8f7b7dde118d..427ea18b7bbf 100644
--- a/python/ray/data/grouped_data.py
+++ b/python/ray/data/grouped_data.py
@@ -1,3 +1,4 @@
+from functools import partial
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union

 from ray.data._internal.aggregate import Count, Max, Mean, Min, Std, Sum
@@ -261,7 +262,10 @@ def wrapped_fn(batch, *args, **kwargs):

         # Change the name of the wrapped function so that users see the name of their
         # function rather than `wrapped_fn` in the progress bar.
-        wrapped_fn.__name__ = fn.__name__
+        if isinstance(fn, partial):
+            wrapped_fn.__name__ = fn.func.__name__
+        else:
+            wrapped_fn.__name__ = fn.__name__

         # Note we set batch_size=None here, so it will use the entire block as a batch,
         # which ensures that each group will be contained within a batch in entirety.

diff --git a/python/ray/data/tests/test_all_to_all.py b/python/ray/data/tests/test_all_to_all.py
index cf0cb8b2b2e7..a6b173383145 100644
--- a/python/ray/data/tests/test_all_to_all.py
+++ b/python/ray/data/tests/test_all_to_all.py
@@ -1167,7 +1167,6 @@ def test_groupby_map_groups_multicolumn(
     ray_start_regular_shared, ds_format, num_parts, use_push_based_shuffle
 ):
     # Test built-in count aggregation
-    print(f"Seeding RNG for test_groupby_arrow_count with: {RANDOM_SEED}")
     random.seed(RANDOM_SEED)
     xs = list(range(100))
     random.shuffle(xs)
@@ -1190,6 +1189,33 @@ def test_groupby_map_groups_multicolumn(
     ]

+def test_groupby_map_groups_with_partial():
+    """
+    The partial function name should show up as
+    +- Sort
+       +- MapBatches(func)
+    """
+    from functools import partial
+
+    def func(x, y):
+        return {f"x_add_{y}": [len(x["id"]) + y]}
+
+    df = pd.DataFrame({"id": list(range(100))})
+    df["key"] = df["id"] % 5
+
+    ds = ray.data.from_pandas(df).groupby("key").map_groups(partial(func, y=5))
+    result = ds.take_all()
+
+    assert result == [
+        {"x_add_5": 25},
+        {"x_add_5": 25},
+        {"x_add_5": 25},
+        {"x_add_5": 25},
+        {"x_add_5": 25},
+    ]
+    assert "MapBatches(func)" in ds.__repr__()
+
+
 def test_random_block_order_schema(ray_start_regular_shared):
     df = pd.DataFrame({"a": np.random.rand(10), "b": np.random.rand(10)})
     ds = ray.data.from_pandas(df).randomize_block_order()

[Data] `GroupedData.map_groups()` doesn't allow partial callables
### What happened + What you expected to happen

In [this PR](https://github.com/ray-project/ray/pull/45310), a new requirement was imposed on the `fn` callable given as input to [`GroupedData.map_groups()`](https://docs.ray.io/en/latest/data/api/doc/ray.data.grouped_data.GroupedData.map_groups.html#ray.data.grouped_data.GroupedData.map_groups): that it have a `__name__` attribute. Unfortunately, callables partially parametrized using [`functools.partial()`](https://docs.python.org/3/library/functools.html#functools.partial) have no such attribute, so passing them into `.map_groups()` raises an error: `AttributeError: 'functools.partial' object has no attribute '__name__'`. This did not happen prior to the linked PR.

It's not a huge deal, but it did cause code to break unexpectedly, and I guess technically is in conflict with the type annotations on this method.

```
In [1]: import functools

In [2]: callable(functools.partial(lambda x: x))
Out[2]: True

In [3]: functools.partial(lambda x: x).__name__
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[3], line 1
----> 1 functools.partial(lambda x: x).__name__

AttributeError: 'functools.partial' object has no attribute '__name__'
```

### Versions / Dependencies

`ray >= 2.21`
PY3.10
macOS 14.4

### Reproduction script

```python
>>> import functools
>>> import ray
>>> ds = ray.data.range(10)
>>> ds.groupby("id").map_groups(lambda x: x)  # this is fine
MapBatches(<lambda>)
+- Sort
   +- Dataset(num_rows=10, schema={id: int64})
>>> ds.groupby("id").map_groups(functools.partial(lambda x: x))  # this errors
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[8], line 1
----> 1 ds.groupby("id").map_groups(functools.partial(lambda x: x))

File ~/.pyenv/versions/3.10.13/envs/ev-detection-py310/lib/python3.10/site-packages/ray/data/grouped_data.py:253, in GroupedData.map_groups(self, fn, compute, batch_format, fn_args, fn_kwargs, fn_constructor_args, fn_constructor_kwargs, num_cpus, num_gpus, concurrency, **ray_remote_args)
    249         yield from apply_udf_to_groups(fn, batch, *args, **kwargs)
    251 # Change the name of the wrapped function so that users see the name of their
    252 # function rather than `wrapped_fn` in the progress bar.
--> 253 wrapped_fn.__name__ = fn.__name__
    255 # Note we set batch_size=None here, so it will use the entire block as a batch,
    256 # which ensures that each group will be contained within a batch in entirety.
    257 return sorted_ds._map_batches_without_batch_size_validation(
    258     wrapped_fn,
    259     batch_size=None,
   (...)
    271     **ray_remote_args,
    272 )

AttributeError: 'functools.partial' object has no attribute '__name__'
```

### Issue Severity

Low: It annoys or frustrates me.


diff --git a/python/ray/dag/compiled_dag_node.py b/python/ray/dag/compiled_dag_node.py
index acec6c2672cd..c964b890f3d4 100644
--- a/python/ray/dag/compiled_dag_node.py
+++ b/python/ray/dag/compiled_dag_node.py
@@ -185,7 +185,7 @@ def do_profile_tasks(
     """
     try:
         for task in tasks:
-            task.prepare()
+            task.prepare(overlap_gpu_communication=overlap_gpu_communication)

         if not hasattr(self, "__ray_adag_events"):
             self.__ray_adag_events = []

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 d1ac1c68063f..1797068e7e2d 100644
--- a/python/ray/dag/tests/experimental/test_torch_tensor_dag.py
+++ b/python/ray/dag/tests/experimental/test_torch_tensor_dag.py
@@ -182,7 +182,11 @@ def test_torch_tensor_as_dag_input(ray_start_regular):

 @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
-def test_torch_tensor_nccl(ray_start_regular):
+@pytest.mark.parametrize("enable_profiling", [False, True])
+@pytest.mark.parametrize("overlap_gpu_communication", [False, True])
+def test_torch_tensor_nccl(
+    ray_start_regular, monkeypatch, enable_profiling, overlap_gpu_communication
+):
     if not USE_GPU:
         pytest.skip("NCCL tests require GPUs")

@@ -190,6 +194,10 @@ def test_torch_tensor_nccl(ray_start_regular):
         sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
     ), "This test requires at least 2 GPUs"

+    monkeypatch.setattr(
+        ray.dag.constants, "RAY_ADAG_ENABLE_PROFILING", enable_profiling
+    )
+
     actor_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)

     sender = actor_cls.remote()
@@ -204,7 +212,9 @@ def test_torch_tensor_nccl(ray_start_regular):
         dag = dag.with_type_hint(TorchTensorType(transport="nccl"))
         dag = receiver.recv.bind(dag)

-    compiled_dag = dag.experimental_compile()
+    compiled_dag = dag.experimental_compile(
+        _overlap_gpu_communication=overlap_gpu_communication
+    )

     # Test that we can pass different shapes and data.
     for i in range(3):

[core][compiled graphs] Failing test: `test_torch_tensor_dag.py::test_torch_tensor_exceptions[static_shape=True, direct_return=True, overlap_gpu_communication=True]`
### What happened + What you expected to happen

The test `ray/python/ray/dag/tests/experimental/test_torch_tensor_dag.py::test_torch_tensor_exceptions[static_shape=True-direct_return=True-overlap_gpu_communication=True]` fails locally.

The bug is probably reading the buffer allocated on the receiver side in NCCL P2P send/recv before the actual data is sent. Currently the receiver's buffer is all zeros so the output is all zeros. When I changed the allocation function to allocate `torch.ones(...) * 100` instead, the actual output becomes `[100, ..., 100]`.

An interesting finding is that when the code executes faster, this test always fails; but when I added a ton of print statements for debugging, it runs more slowly and the test sometimes passes.

Since this test has `overlap_gpu_communication=True`, it is likely related to overlapping GPU communication with computation. My guess is that the actor reading the tensor did not properly wait for the recv event to finish.

I checked out to the commit that most recently modified the test: #47586, as well as the current HEAD of the `ray-project:master` branch, and the test failed in either case.

Below is an example error message:

```python
    @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
    @pytest.mark.parametrize("static_shape", [False, True])
    @pytest.mark.parametrize("direct_return", [False, True])
    @pytest.mark.parametrize("overlap_gpu_communication", [False, True])
    def test_torch_tensor_exceptions(
        ray_start_regular, static_shape, direct_return, overlap_gpu_communication
    ):
        """
        Test exceptions being thrown by a NCCL sending task.
        """
        if not USE_GPU:
            pytest.skip("NCCL tests require GPUs")

        assert (
            sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
        ), "This test requires at least 2 GPUs"

        actor_cls = TorchTensorWorker.options(num_gpus=1)

        sender = actor_cls.remote()
        receiver = actor_cls.remote()

        with InputNode() as inp:
            dag = sender.send_or_raise.bind(
                inp.shape, inp.dtype, inp.value, inp.raise_exception
            )
            dag = dag.with_type_hint(
                TorchTensorType(
                    _static_shape=static_shape,
                    _direct_return=direct_return,
                    transport="nccl",
                )
            )
            dag = receiver.recv.bind(dag)

        compiled_dag = dag.experimental_compile(
            _overlap_gpu_communication=overlap_gpu_communication
        )

        shape = (10,)
        dtype = torch.float16

        for i in range(3):
            i += 1

            ref = compiled_dag.execute(
                shape=shape,
                dtype=dtype,
                value=i,
                raise_exception=False,
            )
            result = ray.get(ref)
>           assert result == (i, shape, dtype)
E           assert (0.0, torch.S...torch.float16) == (2, (10,), torch.float16)
E             At index 0 diff: 0.0 != 2
E             Full diff:
E             - (2, (10,), torch.float16)
E             + (0.0, torch.Size([10]), torch.float16)

test_torch_tensor_dag.py:861: AssertionError
```

### Versions / Dependencies

Newest version of Ray. Python: 3.9.

### Reproduction script

https://github.com/ray-project/ray/blob/master/python/ray/dag/tests/experimental/test_torch_tensor_dag.py#L813

```python
@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
@pytest.mark.parametrize("static_shape", [False, True])
@pytest.mark.parametrize("direct_return", [False, True])
@pytest.mark.parametrize("overlap_gpu_communication", [False, True])
def test_torch_tensor_exceptions(
    ray_start_regular, static_shape, direct_return, overlap_gpu_communication
):
    """
    Test exceptions being thrown by a NCCL sending task.
    """
    if not USE_GPU:
        pytest.skip("NCCL tests require GPUs")

    assert (
        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
    ), "This test requires at least 2 GPUs"

    actor_cls = TorchTensorWorker.options(num_gpus=1)

    sender = actor_cls.remote()
    receiver = actor_cls.remote()

    with InputNode() as inp:
        dag = sender.send_or_raise.bind(
            inp.shape, inp.dtype, inp.value, inp.raise_exception
        )
        dag = dag.with_type_hint(
            TorchTensorType(
                _static_shape=static_shape,
                _direct_return=direct_return,
                transport="nccl",
            )
        )
        dag = receiver.recv.bind(dag)

    compiled_dag = dag.experimental_compile(
        _overlap_gpu_communication=overlap_gpu_communication
    )

    shape = (10,)
    dtype = torch.float16

    for i in range(3):
        i += 1

        ref = compiled_dag.execute(
            shape=shape,
            dtype=dtype,
            value=i,
            raise_exception=False,
        )
        result = ray.get(ref)
        assert result == (i, shape, dtype)

    # Application level exceptions are thrown to the end ray.get
    ref = compiled_dag.execute(
        shape=shape,
        dtype=dtype,
        value=i,
        raise_exception=True,
    )
    if static_shape or direct_return:
        with pytest.raises(RayChannelError):
            # TODO(swang): Ideally return the RuntimeError thrown by the
            # application instead of a generic RayChannelError.
            ray.get(ref)

        with pytest.raises(RayChannelError):
            # If using static_shape=True or direct_return=True, the DAG is not
            # usable after application-level exceptions.
            ref = compiled_dag.execute(
                shape=shape,
                dtype=dtype,
                value=i,
                raise_exception=False,
            )
    else:
        with pytest.raises(RuntimeError):
            ray.get(ref)

        # DAG should still be usable after application-level exceptions.
        ref = compiled_dag.execute(
            shape=shape,
            dtype=dtype,
            value=i,
            raise_exception=False,
        )
        result = ray.get(ref)
        assert result == (i, shape, dtype)
```

### Issue Severity

High: It blocks me from completing my task.

Hi @AndyUB, Thanks for reporting!

This does sound like a bug in the overlap functionality. However, when I reran the test (`python/ray/dag/tests/experimental/test_torch_tensor_dag.py::test_torch_tensor_exceptions[True-True-True-ray_start_regular0]`) 100 times on a node with L4 GPU, I wasn't able to reproduce it once.
Our CI is also a bit flaky right now (fix in progress) so it's not easy to check how often this failed.

In your experience how often this fails? Any way to make it more reproducible?
You mentioned "An interesting finding is that when the code executes faster, this test always fails", how did you make the code run faster? On a better GPU?
btw, I think the issue is likely due to in `_compute()` we only sync on GPU `recv stream`, but not on CPU:

```
    def _compute(
        self,
        overlap_gpu_communication: bool,
        class_handle,
    ) -> bool:
        input_data = self.reset_and_wait_intermediate_future()
```

```
    def reset_and_wait_intermediate_future(self) -> Any:
        future = self._intermediate_future
        self._intermediate_future = None
        return future.wait()
```

```
class GPUFuture(DAGOperationFuture[Any]):
    def wait(self) -> Any:
        """
        Wait for the future on the current CUDA stream and return the result from
        the GPU operation. This operation does not block CPU.
        """
        import cupy as cp

        current_stream = cp.cuda.get_current_stream()
        current_stream.wait_event(self._event)
        return self._buf
```

And the receiver `_compute()` operation runs the the following method (`TorchTensorWorker.recv`), which directly retrieves the item, shape, and dtype from the GPU tensor without waiting.

```
class TorchTensorWorker:
    def recv(self, tensor):
        # Check that tensor got loaded to the correct device.
        assert tensor.device == self.device
        return (tensor[0].item(), tensor.shape, tensor.dtype)
```

To fix this issue, we will probably need to make CPU synchronize on the `recv stream` in `_compute()`.
cc: @stephanie-wang @rkooo567
Re: test repro, you could try to insert a sleep on the recv stream before queuing the recv.

> And the receiver _compute() operation runs the the following method (TorchTensorWorker.recv), which directly retrieves the item, shape, and dtype from the GPU tensor without waiting.
> To fix this issue, we will probably need to make CPU synchronize on the recv stream in _compute().

Not sure that's the whole story, the read of the item requires GPU->CPU movement and is supposed to get queued on the compute stream after syncing on the recv stream. It would be good to check that the read of the item is happening on the expected stream.
> Hi @AndyUB, Thanks for reporting!
>
> This does sound like a bug in the overlap functionality. However, when I reran the test (`python/ray/dag/tests/experimental/test_torch_tensor_dag.py::test_torch_tensor_exceptions[True-True-True-ray_start_regular0]`) 100 times on a node with L4 GPU, I wasn't able to reproduce it once. Our CI is also a bit flaky right now (fix in progress) so it's not easy to check how often this failed.
>
> In your experience how often this fails? Any way to make it more reproducible? You mentioned "An interesting finding is that when the code executes faster, this test always fails", how did you make the code run faster? On a better GPU?

I found the bug. The test failed because I ran with `RAY_ADAG_ENABLE_PROFILING=1`. The bug is in [`do_profile_tasks`, where `task.prepare` does not `overlap_gpu_communication=overlap_gpu_communication`](https://github.com/ray-project/ray/blob/master/python/ray/dag/compiled_dag_node.py#L188). 🤣 After adding that, I ran for 50 times and it passed every time.


diff --git a/python/ray/data/_internal/arrow_block.py b/python/ray/data/_internal/arrow_block.py
index 3ea83b9f04cc..f12f89d8cceb 100644
--- a/python/ray/data/_internal/arrow_block.py
+++ b/python/ray/data/_internal/arrow_block.py
@@ -32,7 +32,7 @@
 )
 from ray.data._internal.row import TableRow
 from ray.data._internal.table_block import TableBlockAccessor, TableBlockBuilder
-from ray.data._internal.util import find_partitions
+from ray.data._internal.util import NULL_SENTINEL, find_partitions
 from ray.data.block import (
     Block,
     BlockAccessor,
@@ -500,7 +500,6 @@ def sort_and_partition(
         table = sort(self._table, sort_key)
         if len(boundaries) == 0:
             return [table]
-
         return find_partitions(table, boundaries, sort_key)

     def combine(self, sort_key: "SortKey", aggs: Tuple["AggregateFn"]) -> Block:
@@ -634,6 +633,11 @@ def key_fn(r):
             else:
                 return (0,)

+        # Replace Nones with NULL_SENTINEL to ensure safe sorting.
+        def key_fn_with_null_sentinel(r):
+            values = key_fn(r)
+            return [NULL_SENTINEL if v is None else v for v in values]
+
         # Handle blocks of different types.
         blocks = TableBlockAccessor.normalize_block_types(blocks, "arrow")

@@ -642,7 +646,7 @@ def key_fn(r):
                 ArrowBlockAccessor(block).iter_rows(public_row_format=False)
                 for block in blocks
             ],
-            key=key_fn,
+            key=key_fn_with_null_sentinel,
         )
         next_row = None
         builder = ArrowBlockBuilder()
diff --git a/python/ray/data/_internal/planner/exchange/sort_task_spec.py b/python/ray/data/_internal/planner/exchange/sort_task_spec.py
index 299e8793774f..827c4a2c7a51 100644
--- a/python/ray/data/_internal/planner/exchange/sort_task_spec.py
+++ b/python/ray/data/_internal/planner/exchange/sort_task_spec.py
@@ -7,6 +7,7 @@
 from ray.data._internal.progress_bar import ProgressBar
 from ray.data._internal.remote_fn import cached_remote_fn
 from ray.data._internal.table_block import TableBlockAccessor
+from ray.data._internal.util import NULL_SENTINEL
 from ray.data.block import Block, BlockAccessor, BlockExecStats, BlockMetadata
 from ray.types import ObjectRef

@@ -23,7 +24,7 @@ def __init__(
         self,
         key: Optional[Union[str, List[str]]] = None,
         descending: Union[bool, List[bool]] = False,
-        boundaries: Optional[list] = None,
+        boundaries: Optional[List[T]] = None,
     ):
         if key is None:
             key = []
@@ -195,7 +196,23 @@ def sample_boundaries(
         samples_table = builder.build()
         samples_dict = BlockAccessor.for_block(samples_table).to_numpy(columns=columns)
         # This zip does the transposition from list of column values to list of tuples.
-        samples_list = sorted(zip(*samples_dict.values()))
+        samples_list = list(zip(*samples_dict.values()))
+
+        def is_na(x):
+            # Check if x is None or NaN. Type casting to np.array first to avoid
+            # isnan failing on strings and other types.
+            if x is None:
+                return True
+            x = np.asarray(x)
+            if np.issubdtype(x.dtype, np.number):
+                return np.isnan(x)
+            return False
+
+        def key_fn_with_nones(sample):
+            return tuple(NULL_SENTINEL if is_na(x) else x for x in sample)
+
+        # Sort the list, but Nones should be NULL_SENTINEL to ensure safe sorting.
+        samples_list = sorted(samples_list, key=key_fn_with_nones)

         # Each boundary corresponds to a quantile of the data.
         quantile_indices = [
diff --git a/python/ray/data/_internal/util.py b/python/ray/data/_internal/util.py
index 9696074fe66d..5e8c921c3733 100644
--- a/python/ray/data/_internal/util.py
+++ b/python/ray/data/_internal/util.py
@@ -55,6 +55,28 @@
 _pyarrow_dataset: LazyModule = None

+class _NullSentinel:
+    """Sentinel value that sorts greater than any other value."""
+
+    def __eq__(self, other):
+        return isinstance(other, _NullSentinel)
+
+    def __lt__(self, other):
+        return False
+
+    def __le__(self, other):
+        return isinstance(other, _NullSentinel)
+
+    def __gt__(self, other):
+        return True
+
+    def __ge__(self, other):
+        return True
+
+
+NULL_SENTINEL = _NullSentinel()
+
+
 def _lazy_import_pyarrow_dataset() -> LazyModule:
     global _pyarrow_dataset
     if _pyarrow_dataset is None:
@@ -723,6 +745,16 @@ def find_partition_index(
         col_vals = table[col_name].to_numpy()[left:right]
         desired_val = desired[i]

+        # Handle null values - replace them with sentinel values
+        if desired_val is None:
+            desired_val = NULL_SENTINEL
+
+        # Replace None/NaN values in col_vals with sentinel
+        null_mask = col_vals == None  # noqa: E711
+        if null_mask.any():
+            col_vals = col_vals.copy()  # Make a copy to avoid modifying original
+            col_vals[null_mask] = NULL_SENTINEL
+
         prevleft = left
         if descending is True:
             left = prevleft + (

diff --git a/python/ray/data/tests/test_all_to_all.py b/python/ray/data/tests/test_all_to_all.py
index a7b6ec823a9c..cf0cb8b2b2e7 100644
--- a/python/ray/data/tests/test_all_to_all.py
+++ b/python/ray/data/tests/test_all_to_all.py
@@ -123,6 +123,65 @@ def test_unique(ray_start_regular_shared):
         assert mock_validate.call_args_list[0].args[0].names == ["b"]

+@pytest.mark.parametrize("batch_format", ["pandas", "pyarrow"])
+def test_unique_with_nulls(ray_start_regular_shared, batch_format):
+    ds = ray.data.from_items([3, 2, 3, 1, 2, 3, None])
+    assert set(ds.unique("item")) == {1, 2, 3, None}
+    assert len(ds.unique("item")) == 4
+
+    ds = ray.data.from_items(
+        [
+            {"a": 1, "b": 1},
+            {"a": 1, "b": 2},
+            {"a": 1, "b": None},
+            {"a": None, "b": 3},
+            {"a": None, "b": 4},
+        ]
+    )
+    assert set(ds.unique("a")) == {1, None}
+    assert len(ds.unique("a")) == 2
+    assert set(ds.unique("b")) == {1, 2, 3, 4, None}
+    assert len(ds.unique("b")) == 5
+
+    # Check with 3 columns
+    df = pd.DataFrame(
+        {
+            "col1": [1, 2, None, 3, None, 3, 2],
+            "col2": [None, 2, 2, 3, None, 3, 2],
+            "col3": [1, None, 2, None, None, None, 2],
+        }
+    )
+    # df["col"].unique() works fine, as expected
+    ds2 = ray.data.from_pandas(df)
+    ds2 = ds2.map_batches(lambda x: x, batch_format=batch_format)
+    assert set(ds2.unique("col1")) == {1, 2, 3, None}
+    assert len(ds2.unique("col1")) == 4
+    assert set(ds2.unique("col2")) == {2, 3, None}
+    assert len(ds2.unique("col2")) == 3
+    assert set(ds2.unique("col3")) == {1, 2, None}
+    assert len(ds2.unique("col3")) == 3
+
+    # Check with 3 columns and different dtypes
+    df = pd.DataFrame(
+        {
+            "col1": [1, 2, None, 3, None, 3, 2],
+            "col2": [None, 2, 2, 3, None, 3, 2],
+            "col3": [1, None, 2, None, None, None, 2],
+        }
+    )
+    df["col1"] = df["col1"].astype("Int64")
+    df["col2"] = df["col2"].astype("Float64")
+    df["col3"] = df["col3"].astype("string")
+    ds3 = ray.data.from_pandas(df)
+    ds3 = ds3.map_batches(lambda x: x, batch_format=batch_format)
+    assert set(ds3.unique("col1")) == {1, 2, 3, None}
+    assert len(ds3.unique("col1")) == 4
+    assert set(ds3.unique("col2")) == {2, 3, None}
+    assert len(ds3.unique("col2")) == 3
+    assert set(ds3.unique("col3")) == {"1.0", "2.0", None}
+    assert len(ds3.unique("col3")) == 3
+
+
 def test_grouped_dataset_repr(ray_start_regular_shared):
     ds = ray.data.from_items([{"key": "spam"}, {"key": "ham"}, {"key": "spam"}])
     assert repr(ds.groupby("key")) == f"GroupedData(dataset={ds!r}, key='key')"
diff --git a/python/ray/data/tests/test_util.py b/python/ray/data/tests/test_util.py
index b66a9bc5804f..e5bfed6154d0 100644
--- a/python/ray/data/tests/test_util.py
+++ b/python/ray/data/tests/test_util.py
@@ -14,6 +14,7 @@
 )
 from ray.data._internal.remote_fn import _make_hashable, cached_remote_fn
 from ray.data._internal.util import (
+    NULL_SENTINEL,
     _check_pyarrow_version,
     _split_list,
     iterate_with_retry,
@@ -35,6 +36,21 @@ def foo():
     assert cpu_only_foo != gpu_only_foo

+def test_null_sentinel():
+    """Check that NULL_SENTINEL sorts greater than any other value."""
+    assert NULL_SENTINEL > 1000
+    assert NULL_SENTINEL > "abc"
+    assert NULL_SENTINEL == NULL_SENTINEL
+    assert NULL_SENTINEL != 1000
+    assert NULL_SENTINEL != "abc"
+    assert not NULL_SENTINEL < 1000
+    assert not NULL_SENTINEL < "abc"
+    assert not NULL_SENTINEL <= 1000
+    assert not NULL_SENTINEL <= "abc"
+    assert NULL_SENTINEL >= 1000
+    assert NULL_SENTINEL >= "abc"
+
+
 def test_make_hashable():
     valid_args = {
         "int": 0,

[Data] Dataset.unique() raises error in case of any null values
### What happened + What you expected to happen

I wanted to get the unique values in a given column of my dataset, but some of the values are null for unavoidable reasons. Calling `Dataset.unique(colname)` on such data raises a TypeError, with differing specifics depending on how the column dtype is specified. This behavior was surprising since the equivalent operation on a `pandas.Series` works just fine, as does getting unique values via Python built-ins.

Here are two versions of type error I got, seemingly from the same line of code:

```
File ~/.pyenv/versions/3.9.18/envs/ev-detection/lib/python3.9/site-packages/ray/data/_internal/planner/exchange/sort_task_spec.py:110, in SortTaskSpec.sample_boundaries(blocks, sort_key, num_reducers)
    107 sample_dict = BlockAccessor.for_block(samples).to_numpy(columns=columns)
    108 # Compute sorted indices of the samples. In np.lexsort last key is the
    109 # primary key hence have to reverse the order.
--> 110 indices = np.lexsort(list(reversed(list(sample_dict.values()))))
    111 # Sort each column by indices, and calculate q-ths quantile items.
    112 # Ignore the 1st item as it's not required for the boundary
    113 for k, v in sample_dict.items():

File <__array_function__ internals>:180, in lexsort(*args, **kwargs)

TypeError: '<' not supported between instances of 'NoneType' and 'int'
```
and

```
File ~/.pyenv/versions/3.9.18/envs/test-env/lib/python3.9/site-packages/ray/data/_internal/planner/exchange/sort_task_spec.py:110, in SortTaskSpec.sample_boundaries(blocks, sort_key, num_reducers)
    107 sample_dict = BlockAccessor.for_block(samples).to_numpy(columns=columns)
    108 # Compute sorted indices of the samples. In np.lexsort last key is the
    109 # primary key hence have to reverse the order.
--> 110 indices = np.lexsort(list(reversed(list(sample_dict.values()))))
    111 # Sort each column by indices, and calculate q-ths quantile items.
    112 # Ignore the 1st item as it's not required for the boundary
    113 for k, v in sample_dict.items():

File <__array_function__ internals>:180, in lexsort(*args, **kwargs)

File missing.pyx:419, in pandas._libs.missing.NAType.__bool__()

TypeError: boolean value of NA is ambiguous
```

### Versions / Dependencies

macOS 14.1
PY 3.9
ray == 2.9.0
pandas == 2.1.0

### Reproduction script

```python
import pandas as pd
import ray.data

items = [1, 2, 3, 2, 3, None]
# set(items) works fine, as expected
ds1 = ray.data.from_items(items)
ds1.unique("item")
# raises TypeError: '<' not supported between instances of 'NoneType' and 'int'

df = pd.DataFrame({"col": [1, 2, 3, None]}, dtype="Int64")
# df["col"].unique() works fine, as expected
ds2 = ray.data.from_pandas(df)
ds2.unique("col")
# raises TypeError: boolean value of NA is ambiguous
```

### Issue Severity

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

Hello burton, I'd like to work on this issue! TIA.
hi @Akshi22 , don't let me get in your way! though it looks like @ujjawal-khare-27 has already submitted a pr to fix this issue. maybe you can help there?
For what it's worth, I just ran into this issue again, only this time in the context of `Dataset.groupby(col)`. It's the same error message, and presumably the same code under the hood. Just a bummer.
Hi, is this issue still open?
If so, I'd like to get started contributing to Ray.io!
I believe the right way to fix this is going to require the underlying Merge operations to be Pyarrow based, instead of Python based (where we currently use a heapq iterator, which doesn't compare NaNs well)


diff --git a/python/ray/autoscaler/v2/scheduler.py b/python/ray/autoscaler/v2/scheduler.py
index 6ac0b0b3a095..3732a6282632 100644
--- a/python/ray/autoscaler/v2/scheduler.py
+++ b/python/ray/autoscaler/v2/scheduler.py
@@ -1568,6 +1568,10 @@ def _enforce_idle_termination(
         Args:
             ctx: The schedule context.
         """
+        count_by_node_type = ctx.get_cluster_shape()
+        node_type_configs = ctx.get_node_type_configs()
+        terminate_nodes_by_type: Dict[NodeType, int] = defaultdict(int)
+
         nodes = ctx.get_nodes()
         s_to_ms = 1000
         for node in nodes:
@@ -1593,13 +1597,32 @@ def _enforce_idle_termination(
                 # Skip it.
                 if node.idle_duration_ms > ctx.get_idle_timeout_s() * s_to_ms:
                     logger.debug(
-                        "Node {}(idle for {} secs) is needed by the cluster resource "
+                        "Node {} (idle for {} secs) is needed by the cluster resource "
                         "constraints, skip idle termination.".format(
                             node.ray_node_id, node.idle_duration_ms / s_to_ms
                         )
                     )
                 continue

+            # Honor the min_worker_nodes setting for the node type.
+            min_count = 0
+            node_type = node.node_type
+            if node_type in node_type_configs:
+                min_count = node_type_configs[node_type].min_worker_nodes
+            if (
+                count_by_node_type.get(node_type, 0)
+                - terminate_nodes_by_type[node_type]
+                <= min_count
+            ):
+                logger.info(
+                    "Node {} (idle for {} secs) belongs to node_type {} and is "
+                    "required by min_worker_nodes, skipping idle termination.".format(
+                        node.ray_node_id, node.idle_duration_ms / s_to_ms, node_type
+                    )
+                )
+                continue
+
+            terminate_nodes_by_type[node.node_type] += 1
             # The node is idle for too long, terminate it.
             node.status = SchedulingNodeStatus.TO_TERMINATE
             node.termination_request = TerminationRequest(

diff --git a/python/ray/autoscaler/v2/tests/test_scheduler.py b/python/ray/autoscaler/v2/tests/test_scheduler.py
index 3920ee9f3a24..e6d6cb71978d 100644
--- a/python/ray/autoscaler/v2/tests/test_scheduler.py
+++ b/python/ray/autoscaler/v2/tests/test_scheduler.py
@@ -1349,6 +1349,91 @@ def test_idle_termination(idle_timeout_s, has_resource_constraints):
         assert len(to_terminate) == 0

+@pytest.mark.parametrize("min_workers", [0, 1])
+def test_idle_termination_with_min_worker(min_workers):
+    """
+    Test that idle nodes are terminated.
+    """
+    idle_timeout_s = 1
+
+    scheduler = ResourceDemandScheduler(event_logger)
+
+    node_type_configs = {
+        "type_cpu": NodeTypeConfig(
+            name="type_cpu",
+            resources={"CPU": 1},
+            min_worker_nodes=min_workers,
+            max_worker_nodes=5,
+            launch_config_hash="hash1",
+        ),
+        "head_node": NodeTypeConfig(
+            name="head_node",
+            resources={"CPU": 0},
+            launch_config_hash="hash2",
+            min_worker_nodes=0,
+            max_worker_nodes=1,
+        ),
+    }
+
+    idle_time_s = 5
+    constraints = []
+
+    request = sched_request(
+        node_type_configs=node_type_configs,
+        instances=[
+            make_autoscaler_instance(
+                im_instance=Instance(
+                    instance_id="i-1",
+                    instance_type="type_cpu",
+                    status=Instance.RAY_RUNNING,
+                    launch_config_hash="hash1",
+                    node_id="r-1",
+                ),
+                ray_node=NodeState(
+                    ray_node_type_name="type_cpu",
+                    node_id=b"r-1",
+                    available_resources={"CPU": 1},
+                    total_resources={"CPU": 1},
+                    idle_duration_ms=idle_time_s * 1000,
+                    status=NodeStatus.IDLE,
+                ),
+                cloud_instance_id="c-1",
+            ),
+            make_autoscaler_instance(
+                im_instance=Instance(
+                    instance_id="i-2",
+                    instance_type="head_node",
+                    status=Instance.RAY_RUNNING,
+                    launch_config_hash="hash2",
+                    node_kind=NodeKind.HEAD,
+                    node_id="r-2",
+                ),
+                ray_node=NodeState(
+                    ray_node_type_name="head_node",
+                    node_id=b"r-2",
+                    available_resources={"CPU": 0},
+                    total_resources={"CPU": 0},
+                    idle_duration_ms=999 * 1000,  # idle
+                    status=NodeStatus.IDLE,
+                ),
+                cloud_instance_id="c-2",
+            ),
+        ],
+        idle_timeout_s=idle_timeout_s,
+        cluster_resource_constraints=constraints,
+    )
+
+    reply = scheduler.schedule(request)
+    _, to_terminate = _launch_and_terminate(reply)
+    assert idle_timeout_s <= idle_time_s
+    if min_workers == 0:
+        assert len(to_terminate) == 1
+        assert to_terminate == [("i-1", "r-1", TerminationRequest.Cause.IDLE)]
+    else:
+        assert min_workers > 0
+        assert len(to_terminate) == 0
+
+
 def test_gang_scheduling():
     """
     Test that gang scheduling works.

[Cluster][Autoscaler-v2]-Autoscaler v2 does not honor minReplicas/replicas count of the worker nodes and constantly terminates after idletimeout
### What happened + What you expected to happen

**The bug:**
When the ray cluster is idle, ray autoscaler v2 constantly tries to terminate worker nodes ignoring the set minimum worker nodes count. This causes the ray worker nodes count to go below the minimum set count from the kuberay chart.
Example: consider a ray cluster provisioned in kubernetes using the kuberay chart 1.1.0 with autoscaler v2 enabled and minimum worker nodes setting of 3. When idle, autoscaler terminates the worker nodes after idle timeout seconds causing the active worker node count to fall below 3 (1 or 2 or 0 sometimes for a brief period)

Due to this constant terminate and recreate cycle, sometimes we see the Actors failing as follows:
```
ray.exceptions.ActorDiedError: The actor died unexpectedly before finishing this task.
    class_name: DB
    actor_id: c78dd85202773ed8f736b21c72000000
    namespace: 81a40ef8-077d-4f65-be29-4d4077c23a85
The actor died because its node has died. Node Id: xxxxxx
   the actor's node was terminated: Termination of node that's idle for 244.98 seconds.
The actor never ran - it was cancelled before it started running.
```

**Expected behavior:**
Autoscaler should honor the minimum worker nodes count setting and should not terminate the nodes when the termination causes node count to fall below the set value. This was properly happening in autoscaler v1 and this issue was introduced after upgrading to autoscaler v2.

**Autoscaler logs:**
```
2024-09-09 09:11:41,933	INFO cloud_provider.py:448 -- Listing pods for RayCluster xxxxxxxtest-kuberay in namespace xxxxxxx-test at pods resource version >= 1198596159.
2024-09-09 09:11:41,946 - INFO - Fetched pod data at resource version 1198599914.
2024-09-09 09:11:41,946	INFO cloud_provider.py:466 -- Fetched pod data at resource version 1198599914.
2024-09-09 09:11:41,949 - INFO - Removing 2 nodes of type workergroup (idle).
2024-09-09 09:11:41,949	INFO event_logger.py:76 -- Removing 2 nodes of type workergroup (idle).
2024-09-09 09:11:41,950 - INFO - Update instance RAY_RUNNING->RAY_STOP_REQUESTED (id=c6d6af81-80b8-46b8-8694-24fb726c02dd, type=workergroup, cloud_instance_id=xxxxxxxtest-kuberay-worker-workergroup-7kzpp, ray_id=0f379b4e4c7957d692711e1f505c25425befe61a08c54ee80fc407f2): draining ray: idle for 242.003 secs > timeout=240.0 secs
2024-09-09 09:11:41,950	INFO instance_manager.py:262 -- Update instance RAY_RUNNING->RAY_STOP_REQUESTED (id=c6d6af81-80b8-46b8-8694-24fb726c02dd, type=workergroup, cloud_instance_id=xxxxxxxtest-kuberay-worker-workergroup-7kzpp, ray_id=0f379b4e4c7957d692711e1f505c25425befe61a08c54ee80fc407f2): draining ray: idle for 242.003 secs > timeout=240.0 secs
2024-09-09 09:11:41,950 - INFO - Update instance RAY_RUNNING->RAY_STOP_REQUESTED (id=a93be1d8-dd1e-489c-8f50-6b03087e5280, type=workergroup, cloud_instance_id=xxxxxxxtest-kuberay-worker-workergroup-xn279, ray_id=1bfae5adba89e5437380731e5d53f34b31862e9e46e554512e52e914): draining ray: idle for 242.986 secs > timeout=240.0 secs
2024-09-09 09:11:41,950	INFO instance_manager.py:262 -- Update instance RAY_RUNNING->RAY_STOP_REQUESTED (id=a93be1d8-dd1e-489c-8f50-6b03087e5280, type=workergroup, cloud_instance_id=xxxxxxxtest-kuberay-worker-workergroup-xn279, ray_id=1bfae5adba89e5437380731e5d53f34b31862e9e46e554512e52e914): draining ray: idle for 242.986 secs > timeout=240.0 secs
2024-09-09 09:11:41,955 - INFO - Drained ray on 0f379b4e4c7957d692711e1f505c25425befe61a08c54ee80fc407f2(success=True, msg=)
2024-09-09 09:11:41,955	INFO ray_stopper.py:116 -- Drained ray on 0f379b4e4c7957d692711e1f505c25425befe61a08c54ee80fc407f2(success=True, msg=)
2024-09-09 09:11:41,959 - INFO - Drained ray on 1bfae5adba89e5437380731e5d53f34b31862e9e46e554512e52e914(success=True, msg=)
2024-09-09 09:11:41,959	INFO ray_stopper.py:116 -- Drained ray on 1bfae5adba89e5437380731e5d53f34b31862e9e46e554512e52e914(success=True, msg=)
2024-09-09 09:11:46,986 - INFO - Calculating hashes for file mounts and ray commands.
2024-09-09 09:11:46,986	INFO config.py:182 -- Calculating hashes for file mounts and ray commands.
2024-09-09 09:11:47,015 - INFO - Listing pods for RayCluster xxxxxxxtest-kuberay in namespace xxxxxxx-test at pods resource version >= 1198596159.
2024-09-09 09:11:47,015	INFO cloud_provider.py:448 -- Listing pods for RayCluster xxxxxxxtest-kuberay in namespace xxxxxxx-test at pods resource version >= 1198596159.
2024-09-09 09:11:47,028 - INFO - Fetched pod data at resource version 1198600006.
2024-09-09 09:11:47,028	INFO cloud_provider.py:466 -- Fetched pod data at resource version 1198600006.
2024-09-09 09:11:47,030 - INFO - Update instance RAY_STOP_REQUESTED->RAY_STOPPED (id=c6d6af81-80b8-46b8-8694-24fb726c02dd, type=workergroup, cloud_instance_id=xxxxxxxtest-kuberay-worker-workergroup-7kzpp, ray_id=0f379b4e4c7957d692711e1f505c25425befe61a08c54ee80fc407f2): ray node 0f379b4e4c7957d692711e1f505c25425befe61a08c54ee80fc407f2 is DEAD
2024-09-09 09:11:47,030	INFO instance_manager.py:262 -- Update instance RAY_STOP_REQUESTED->RAY_STOPPED (id=c6d6af81-80b8-46b8-8694-24fb726c02dd, type=workergroup, cloud_instance_id=xxxxxxxtest-kuberay-worker-workergroup-7kzpp, ray_id=0f379b4e4c7957d692711e1f505c25425befe61a08c54ee80fc407f2): ray node 0f379b4e4c7957d692711e1f505c25425befe61a08c54ee80fc407f2 is DEAD
2024-09-09 09:11:47,030 - INFO - Update instance RAY_STOP_REQUESTED->RAY_STOPPED (id=a93be1d8-dd1e-489c-8f50-6b03087e5280, type=workergroup, cloud_instance_id=xxxxxxxtest-kuberay-worker-workergroup-xn279, ray_id=1bfae5adba89e5437380731e5d53f34b31862e9e46e554512e52e914): ray node 1bfae5adba89e5437380731e5d53f34b31862e9e46e554512e52e914 is DEAD
2024-09-09 09:11:47,030	INFO instance_manager.py:262 -- Update instance RAY_STOP_REQUESTED->RAY_STOPPED (id=a93be1d8-dd1e-489c-8f50-6b03087e5280, type=workergroup, cloud_instance_id=xxxxxxxtest-kuberay-worker-workergroup-xn279, ray_id=1bfae5adba89e5437380731e5d53f34b31862e9e46e554512e52e914): ray node 1bfae5adba89e5437380731e5d53f34b31862e9e46e554512e52e914 is DEAD
2024-09-09 09:11:47,031 - INFO - Adding 2 nodes to satisfy min count for node type: workergroup.
2024-09-09 09:11:47,031	INFO scheduler.py:1148 -- Adding 2 nodes to satisfy min count for node type: workergroup.
2024-09-09 09:11:47,032 - INFO - Adding 2 node(s) of type workergroup.
2024-09-09 09:11:47,032	INFO event_logger.py:56 -- Adding 2 node(s) of type workergroup.
```

### Versions / Dependencies

kuberay: v1.1.0
ray-cluster helm chart: v1.1.0
ray version: 2.34.0

### Reproduction script

You can use the base kuberay helm chart to deploy a ray cluster with the following changes:
```
head:
  enableInTreeAutoscaling: true
  autoscalerOptions:
    upscalingMode: Conservative
    idleTimeoutSeconds: 240
  containerEnv:
    - name: RAY_enable_autoscaler_v2
      value: "1"

worker:
  replicas: 3
  minReplicas: 3
  maxReplicas: 10
```

### Issue Severity

High: It blocks me from completing my task.

@jjyao I'm putting this as p0.5 since it's on KubeRay and ASv2 is supposed to work... we could downgrade since v2 isn't quite GA yet...
> @jjyao I'm putting this as p0.5 since it's on KubeRay and ASv2 is supposed to work... we could downgrade since v2 isn't quite GA yet...

@anyscalesam  Wouldn't this be related to ray itself since the autoscaler container is nothing but the ray image?
Yes, we can downgrade autoscaler but the previous version has its own issues and terminates ray nodes even when there are running tasks.


diff --git a/python/ray/_private/worker.py b/python/ray/_private/worker.py
index 3a895356fed8..f7bbd1555ba2 100644
--- a/python/ray/_private/worker.py
+++ b/python/ray/_private/worker.py
@@ -853,6 +853,7 @@ def get_objects(
         object_refs: list,
         timeout: Optional[float] = None,
         return_exceptions: bool = False,
+        skip_deserialization: bool = False,
     ):
         """Get the values in the object store associated with the IDs.

@@ -869,8 +870,10 @@ def get_objects(
                 Exception object, whether to return them as values in the
                 returned list. If False, then the first found exception will be
                 raised.
+            skip_deserialization: If true, only the buffer will be released and
+                the object associated with the buffer will not be deserailized.
         Returns:
-            list: List of deserialized objects
+            list: List of deserialized objects or None if skip_deserialization is True.
             bytes: UUID of the debugger breakpoint we should drop
                 into or b"" if there is no breakpoint.
         """
@@ -903,6 +906,9 @@ def get_objects(
                     debugger_breakpoint = metadata_fields[1][
                         len(ray_constants.OBJECT_METADATA_DEBUG_PREFIX) :
                     ]
+        if skip_deserialization:
+            return None, debugger_breakpoint
+
         values = self.deserialize_objects(data_metadata_pairs, object_refs)
         if not return_exceptions:
             # Raise exceptions instead of returning them to the user.
diff --git a/python/ray/dag/compiled_dag_node.py b/python/ray/dag/compiled_dag_node.py
index 522e0779f2eb..1a4d7a7abe08 100644
--- a/python/ray/dag/compiled_dag_node.py
+++ b/python/ray/dag/compiled_dag_node.py
@@ -2056,6 +2056,21 @@ def _get_execution_results(
                 del self._result_buffer[execution_index]
         return result

+    def release_output_channel_buffers(self, execution_index: int):
+        from ray.dag import DAGContext
+
+        ctx = DAGContext.get_current()
+        timeout = ctx.retrieval_timeout
+
+        while self._max_finished_execution_index < execution_index:
+            self.increment_max_finished_execution_index()
+            start_time = time.monotonic()
+            self._dag_output_fetcher.release_channel_buffers(timeout)
+
+            if timeout != -1:
+                timeout -= time.monotonic() - start_time
+                timeout = max(timeout, 0)
+
     def _execute_until(
         self,
         execution_index: int,
diff --git a/python/ray/experimental/channel/common.py b/python/ray/experimental/channel/common.py
index d5f42178e1c9..9545c70c1883 100644
--- a/python/ray/experimental/channel/common.py
+++ b/python/ray/experimental/channel/common.py
@@ -342,6 +342,14 @@ def _read_list(self, timeout: Optional[float] = None) -> List[Any]:
                 timeout = max(timeout, 0)
         return results

+    def release_channel_buffers(self, timeout: Optional[float] = None) -> None:
+        for c in self._input_channels:
+            start_time = time.monotonic()
+            c.release_buffer(timeout)
+            if timeout is not None:
+                timeout -= time.monotonic() - start_time
+                timeout = max(timeout, 0)
+

 @DeveloperAPI
 class AwaitableBackgroundReader(ReaderInterface):
diff --git a/python/ray/experimental/channel/shared_memory_channel.py b/python/ray/experimental/channel/shared_memory_channel.py
index d9ada000df10..773f85a69f75 100644
--- a/python/ray/experimental/channel/shared_memory_channel.py
+++ b/python/ray/experimental/channel/shared_memory_channel.py
@@ -503,6 +503,18 @@ def read(self, timeout: Optional[float] = None) -> Any:

         return ret

+    def release_buffer(self, timeout: Optional[float] = None) -> None:
+        assert (
+            timeout is None or timeout >= 0 or timeout == -1
+        ), "Timeout must be non-negative or -1."
+        self.ensure_registered_as_reader()
+        self._worker.get_objects(
+            [self._local_reader_ref],
+            timeout=timeout,
+            return_exceptions=True,
+            skip_deserialization=True,
+        )
+
     def close(self) -> None:
         """
         Close this channel by setting the error bit on both the writer_ref and the
@@ -604,6 +616,21 @@ def read(self, timeout: Optional[float] = None) -> Any:
         self._next_read_index %= self._num_shm_buffers
         return output

+    def release_buffer(self, timeout: Optional[float] = None):
+        """Release the native buffer of the channel to allow the buffer to be reused for
+        future data.
+
+        If the next buffer is available, it returns immediately. If the next
+        buffer is not written by an upstream producer, it blocks until a buffer is
+        available to be released. If a buffer is not available within timeout, it raises
+        RayChannelTimeoutError.
+        """
+        # A single channel is not supposed to read and write at the same time.
+        assert self._next_write_index == 0
+        self._buffers[self._next_read_index].release_buffer(timeout)
+        self._next_read_index += 1
+        self._next_read_index %= self._num_shm_buffers
+
     def close(self) -> None:
         for buffer in self._buffers:
             buffer.close()
@@ -732,16 +759,23 @@ def write(self, value: Any, timeout: Optional[float] = None) -> None:

     def read(self, timeout: Optional[float] = None) -> Any:
         self.ensure_registered_as_reader()
+        return self._channel_dict[self._resolve_actor_id()].read(timeout)
+
+    def release_buffer(self, timeout: Optional[float] = None):
+        self.ensure_registered_as_reader()
+        self._channel_dict[self._resolve_actor_id()].release_buffer(timeout)
+
+    def _resolve_actor_id(self) -> str:
         actor_id = ray.get_runtime_context().get_actor_id()
-        # if actor_id is None, read was called by the driver
-        # if the driver is an actor, driver_actor_id will be set to that actor id
+        # If actor_id is None, read was called by the driver
+        # If the driver is an actor, driver_actor_id will be set to that actor id
         if actor_id is None or actor_id == self._driver_actor_id:
             # Use the actor ID of the DAGDriverProxyActor.
             # The proxy actor is always the first actor in the reader_and_node_list.
             assert len(self._reader_and_node_list) >= 1
             driver_proxy_actor = self._reader_and_node_list[0][0]
             actor_id = self._get_actor_id(driver_proxy_actor)
-        return self._channel_dict[actor_id].read(timeout)
+        return actor_id

     def close(self) -> None:
         for channel in self._channels:
diff --git a/python/ray/experimental/compiled_dag_ref.py b/python/ray/experimental/compiled_dag_ref.py
index da53c68d8f8a..4032f8bccc2f 100644
--- a/python/ray/experimental/compiled_dag_ref.py
+++ b/python/ray/experimental/compiled_dag_ref.py
@@ -88,9 +88,11 @@ def __del__(self):
         if self._dag.is_teardown:
             return

-        # If not yet, get the result and discard to avoid execution result leak.
+        # If not yet, release native buffers to avoid execution result leak. Note that
+        # we skip python-based deserialization as the values stored in the buffers are
+        # not used.
         if not self._ray_get_called:
-            self.get()
+            self._dag.release_output_channel_buffers(self._execution_index)

     def get(self, timeout: Optional[float] = None):
         if self._ray_get_called:

diff --git a/python/ray/dag/tests/experimental/test_accelerated_dag.py b/python/ray/dag/tests/experimental/test_accelerated_dag.py
index c764442a6c44..e457d8e9b95b 100644
--- a/python/ray/dag/tests/experimental/test_accelerated_dag.py
+++ b/python/ray/dag/tests/experimental/test_accelerated_dag.py
@@ -170,6 +170,37 @@ def test_basic(ray_start_regular):
         del result

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

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

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

### Versions / Dependencies

ray master

### Reproduction script

```
a = Actor.remote(0)
with InputNode() as inp:
    dag = a.inc.bind(inp)

compiled_dag = dag.experimental_compile()
ref = compiled_dag.execute(1)
compiled_dag.teardown()

# ref.get() is called upon ref's deallocation -- deserialization still happens
```

### Issue Severity

None

@jeffreyjeffreywang would you be interested in taking this?
@rkooo567 Yup, definitely!

Hey @stephanie-wang @ruisearch42, I just wanted to clarify a few things before I proceed with this issue. You've suggested to release the value of a CompiledDAGRef if `ray.get()` isn't called by the user (https://github.com/ray-project/ray/pull/45951#issuecomment-2168530073). **Could you please help me understand how the following code avoids execution result leak?** https://github.com/ray-project/ray/blob/44dd9a7abd0a153dd11167c8d48cde9b0a9ecfb9/python/ray/experimental/compiled_dag_ref.py#L87-L89

With this code, another issue arises -- attempting to load a python library during program exit (when CompiledDAGRef is destructed) will fail. Please refer to https://github.com/ray-project/ray/pull/47305#discussion_r1738029669 for more context.
```
foo = Foo.remote()
bar = Bar.remote()

with InputNode() as inp:
    dag = MultiOutputNode([foo.increment.bind(inp), bar.decrement.bind(inp)])

dag = dag.experimental_compile()

ref1 = dag.execute(1)
ref2 = dag.execute(1)

assert ref1.get() == [1, -1]

dag.teardown()
# Upon destruction, the DAG will be executed until the latest index.
# However, we attempt to import DAGContext during program exit.
```

I'm thinking about removing the custom destructor entirely but wanted to understand the implications before doing so.

> With this code, another issue arises -- attempting to load a python library during program exit (when CompiledDAGRef is destructed) will fail. Please refer to https://github.com/ray-project/ray/pull/47305#discussion_r1738029669 for more context.

I think this is because python cannot guarantee all modules exist when `__del__` is called. So I agreed doing clean up this way in a destructor is a bad idea, and it is something we should clean up.

> Could you please help me understand how the following code avoids execution result leak?

With this code, the python object is retrieved and then immediately goes out of scope. If there are any native buffers underneath, they will also be released.
> With this code, the python object is retrieved and then immediately goes out of scope. If there are any native buffers underneath, they will also be released.

Thank you @ruisearch42, could you give me an example/repro when native buffers are used and therefore this destruction is necessary? I'd like to measure whether the deserialization is necessary.

If deserialization is necessary, we still need to solve the module import issue. We might want to move the deserialization (stepping through the remaining steps in the DAG) to `teardown()`.

This is a duplicate of #46909. Will close both bugs once this issue is addressed.
Thanks, @jeffreyjeffreywang for the great questions! The deserialization is necessary because the native buffer is reused for future data. If the reader does not explicitly read and release the buffer, then the buffer cannot be reused for future values. You can reproduce it by returning a numpy array as the DAG output; since numpy arrays are zero-copy, the buffer will be held until the np array in python goes out of scope.

Note that you do not need to deserialize the data in order to release the buffer. We just need to make sure to call the [`ReadAcquire`](https://github.com/ray-project/ray/blob/35b453445b49ad286c08099953c000c3942acf6a/src/ray/core_worker/experimental_mutable_object_manager.h#L180) and `ReadRelease` methods on the buffer (but skip the python-based deserialization).

We do a similar custom destructor for when ObjectRefs and actors go out of scope, so I think you can reuse a similar codepath to avoid the destruction ordering problem, see [here](https://github.com/ray-project/ray/blob/35b453445b49ad286c08099953c000c3942acf6a/python/ray/actor.py#L1380).
Thank you, Stephanie, for the thorough explanation. I'll dig a bit deeper and publish a PR.
Thanks @jeffreyjeffreywang ! Btw, are you in OSS ray slack? We have regular sync up, and you are more than welcome to join!
@rkooo567 Yeah, I just joined couple days ago. Thanks for inviting, I'll keep an eye on the sync up next time and hopefully I'll be able to join! 😄
@anyscalesam can you make sure @jeffreyjeffreywang is invited to next sync?! Thank you!