Repository: ray-project/ray
Pull request #51032: [Compiled Graph] Enhance Compile Graph with Multi-Device Support
Author: login: hipudding
Labels: name: core
description: Issues that should be addressed in Ray Core
name: go
description: add ONLY when ready to merge, run all tests
name: community-contribution
description: Contributed by the community
Created: 2025-03-03T11:47:28Z | Merged: 2025-05-22T20:10:28Z

<!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?
This PR improves multi-device support in Compile Graph, which significantly reduces Tensor transmission latency by utilizing out-of-band communication. Currently, this feature only supports CUDA’s NCCL. Since Ray already supports multiple accelerators, it is necessary to extend Compile Graph to support multi-device as well.

This PR mainly introduces two key changes:
	1.	Removed dependency on cupy.cuda.ExternalStream – Since this library only supports CUDA devices, we replaced it with a more general stream context manager to accommodate various accelerators. The new implementation uses torch.{device}.StreamContext.
	2.	Replaced hardcoded torch.cuda.xxx calls with AcceleratorRuntime – This allows automatic detection of the accelerator type and invokes the appropriate device-specific functions.

### How to add a new backend for CG? here's an example for Ascend NPU:
```python
import ray
import torch
import torch_npu
from ray.dag import InputNode
# implement customer Communicator class
from ray.experimental.channel.hccl_group import _HcclGroup
from ray.experimental.channel.accelerator_context import register_accelerator_context

@ray.remote
class TorchTensorWorker:
    def __init__(self):
        self.device = torch.device('npu:0')
        torch.npu.set_device(self.device)

    def send(self, shape, dtype, value: int):
        return torch.ones(shape, dtype=dtype, device=self.device) * value

    def recv(self, tensor):
        return (tensor[0].item(), tensor.shape, tensor.dtype)

# global register accelerator context
register_accelerator_context('npu', _HcclGroup)

actor_cls = TorchTensorWorker.options(num_cpus=0, resources={'NPU': 1})

sender = actor_cls.remote()
receiver = actor_cls.remote()

with InputNode() as inp:
    dag = sender.send.bind(inp.shape, inp.dtype, inp[0])
    dag = dag.with_tensor_transport(transport='nccl')
    dag = receiver.recv.bind(dag)

shape = (10,)
dtype = torch.float16

compiled_dag = dag.experimental_compile()
for i in range(3):
    ref = compiled_dag.execute(i, shape=shape, dtype=dtype)
    assert ray.get(ref) == (i, shape, dtype)

print("Success")
```


<!-- Please give a short summary of the change and the problem this solves. -->

## Related issue number

<!-- For example: "Closes #1234" -->
This PR is the main part of Task 2 in #51574 
It would better to set the function name more general, such as changing requires_nccl to require_communicator. This is implemented in #51061.

## Checks

- [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR.
- [x] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for https://docs.ray.io/en/master/.
    - [ ] I've added any new APIs to the API Reference. For example, if I added a
           method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [x] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [x] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(


Linked issue #1234: Ray fails to serialize self-reference objects
<!--
General questions should be asked on the mailing list ray-dev@googlegroups.com.

Before submitting an issue, please fill out the following form.
-->

### System information
- **Ray installed from (source or binary)**: pip
- **Ray version**: 0.2.2
- **Python version**: 3.6.2

<!--
You can obtain the Ray version with

python -c "import ray; print(ray.__version__)"
-->

### Describe the problem

Ray fails to serialize self-reference objects (for example, Graph objects in networkx).

I think it is because ray always tries to use pyarrow first and does not catch `pyarrow.lib.ArrowNotImplementedError`, see

https://github.com/ray-project/ray/blob/e0360eb4298371e308cd266dd512befaa457ce24/python/ray/worker.py#L285-L289

After catching `pyarrow.lib.ArrowNotImplementedError`, we **should not** use `use_dict=True` as a workaround, because it will cause endless loop. A correct approach may be:

```python
            except (pyarrow.SerializationCallbackError, pyarrow.lib.ArrowNotImplementedError) as e:
                try:
                    if isinstance(e, pyarrow.lib.ArrowNotImplementedError):
                        e.example_object = value
                        raise e  # redirect to use cloudpickle
```

### Source code / logs
<!-- Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. -->

```python
class Graph:
    def __init__(self):
        self.g = self

G = Graph()
ray.put(G)  # --> pyarrow.lib.ArrowNotImplementedError: This object exceeds the maximum recursion depth. It may contain itself recursively.

# another example

import networkx as nx
G = nx.Graph()
    
G.add_edges_from([(1, 2), (1, 3)])
G.add_node(1)
G.add_edge(1, 2)
G.add_node("spam")  # adds node "spam"
G.add_nodes_from("spam")  # adds 4 nodes: 's', 'p', 'a', 'm'
G.add_edge(3, 'm')
ray.put(G)  # --> pyarrow.lib.ArrowNotImplementedError: This object exceeds the maximum recursion depth. It may contain itself recursively.
```

@mitar 
Linked issue comments: user: login: robertnishihara
created_at: 2017-11-20T05:34:11Z
body: Can you try `ray.register_custom_serializer`? The following works for me.

```python
import ray
ray.init()

class Graph:
    def __init__(self):
        self.g = self

ray.register_custom_serializer(Graph, use_pickle=True)

G = Graph()
ray.put(G)
```

This is closely related to #319 and https://issues.apache.org/jira/browse/ARROW-1382.

A side comment. The original code worked for me in Python 2 because in Python 2 `Graph` is an old-style class and so we automatically fall back to Pickle anyway I think.
author_association: COLLABORATOR
user: login: mitar
created_at: 2017-11-20T06:02:44Z
body: Hm, so ideally we would like to serialize networkx graphs. Because they can be quite large, I am not sure if pickling is a good approach.
author_association: MEMBER
user: login: robertnishihara
created_at: 2017-11-20T06:32:14Z
body: Custom serializers/deserializers can be registered with the same approach. Not sure what the right one would be in this case, but just as a simple example, you could do something like

```python
import numpy as np
import ray

ray.init()

class Graph:
    def __init__(self, big_array):
        self.g = self
        self.big_array = big_array

def custom_graph_serializer(obj):
    return obj.big_array

def custom_graph_deserializer(serialized_obj):
    return Graph(serialized_obj)

ray.register_custom_serializer(Graph,
                               serializer=custom_graph_serializer,
                               deserializer=custom_graph_deserializer)

G = Graph(np.ones(100))
ray.put(G)
```
author_association: COLLABORATOR
user: login: edoakes
created_at: 2020-03-05T23:11:03Z
body: Stale - please open new issue if still relevant
author_association: COLLABORATOR

Diff:
diff --git a/python/ray/dag/compiled_dag_node.py b/python/ray/dag/compiled_dag_node.py
index 99c3ad35db03..8fde385cbc37 100644
--- a/python/ray/dag/compiled_dag_node.py
+++ b/python/ray/dag/compiled_dag_node.py
@@ -4,7 +4,6 @@
 from contextlib import nullcontext
 from dataclasses import dataclass, asdict
 from typing import (
-    TYPE_CHECKING,
     Any,
     Dict,
     List,
@@ -83,8 +82,7 @@
 
 from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
 
-if TYPE_CHECKING:
-    import cupy as cp
+from ray.experimental.channel.accelerator_context import AcceleratorContext
 
 logger = logging.getLogger(__name__)
 
@@ -371,17 +369,19 @@ def _device_context_manager():
         return nullcontext()
 
     import torch
+    from ray.experimental.channel.accelerator_context import AcceleratorContext
 
     device = ChannelContext.get_current().torch_device
 
-    if device.type == "cuda" and torch.cuda.is_available():
+    if device.type == "cuda" and not torch.cuda.is_available():
         # In the case of mocked NCCL, we may get a device with type "cuda"
         # but CUDA is not available. We return nullcontext() in that case,
         # otherwise torch raises a runtime error if the cuda device context
         # manager is used.
         # TODO(rui): consider better mocking NCCL to support device context.
-        return torch.cuda.device(device)
-    return nullcontext()
+        return nullcontext()
+
+    return AcceleratorContext.get().get_device_context(device)
 
 
 @DeveloperAPI
@@ -590,8 +590,10 @@ def prepare(self, overlap_gpu_communication: bool = False):
         self.input_reader.start()
         self.output_writer.start()
 
-        self._send_stream: Union["cp.cuda.Stream", nullcontext] = nullcontext()
-        self._recv_stream: Union["cp.cuda.Stream", nullcontext] = nullcontext()
+        # Stream context type are different between different accelerators.
+        # Type hint is not applicable here.
+        self._send_stream = nullcontext()
+        self._recv_stream = nullcontext()
         if not overlap_gpu_communication:
             return
 
@@ -674,7 +676,8 @@ def _read(self, overlap_gpu_communication: bool) -> bool:
             # a GPUFuture so that this read operation (communication) can
             # be overlapped with computation.
             self.wrap_and_set_intermediate_future(
-                input_data, wrap_in_gpu_future=overlap_gpu_communication
+                input_data,
+                wrap_in_gpu_future=overlap_gpu_communication,
             )
         except RayChannelError:
             # Channel closed. Exit the loop.
@@ -1321,6 +1324,10 @@ def _init_communicators(self) -> None:
             for type_hint in type_hints:
                 type_hint.set_communicator_id(communicator_id)
 
+        # Second, get registered accelerator context if any.
+        accelerator_module_name = AcceleratorContext.get().module_name
+        accelerator_communicator_cls = AcceleratorContext.get().communicator_cls
+
         # Then, create communicators for collective operations.
         # Reuse an already created communicator for the same set of actors.
         for collective_op in self._collective_ops_with_unresolved_communicators:
@@ -1337,6 +1344,8 @@ def _init_communicators(self) -> None:
                     list(actors),
                     None,
                     self._overlap_gpu_communication,
+                    accelerator_module_name,
+                    accelerator_communicator_cls,
                 )
                 self._actors_to_created_communicator_id[actors] = communicator_id
             collective_op.type_hint.set_communicator_id(communicator_id)
@@ -1358,6 +1367,8 @@ def _init_communicators(self) -> None:
                     list(self._p2p_actors_with_unresolved_communicators),
                     None,
                     self._overlap_gpu_communication,
+                    accelerator_module_name,
+                    accelerator_communicator_cls,
                 )
             for dag_node in self._p2p_dag_nodes_with_unresolved_communicators:
                 dag_node.type_hint.set_communicator_id(p2p_communicator_id)
diff --git a/python/ray/dag/dag_node.py b/python/ray/dag/dag_node.py
index 1587e7976bc2..b2e9fc3f9e0e 100644
--- a/python/ray/dag/dag_node.py
+++ b/python/ray/dag/dag_node.py
@@ -172,9 +172,9 @@ def with_tensor_transport(
         try:
             device = Device(device)
         except ValueError:
+            valid_devices = ", ".join(f"'{d.value}'" for d in Device)
             raise ValueError(
-                f"Invalid device '{device}'. "
-                "Valid options are: 'default', 'cpu', 'gpu', 'cuda'."
+                f"Invalid device '{device}'. Valid options are: {valid_devices}."
             )
         if transport == "auto":
             self._type_hint = AutoTransportType(
diff --git a/python/ray/dag/dag_operation_future.py b/python/ray/dag/dag_operation_future.py
index cc7fd174aa2b..acfc83d7c1d6 100644
--- a/python/ray/dag/dag_operation_future.py
+++ b/python/ray/dag/dag_operation_future.py
@@ -1,11 +1,9 @@
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar, Dict
+from typing import Any, Generic, TypeVar, Dict
 from ray.util.annotations import DeveloperAPI
+from ray.experimental.channel.accelerator_context import AcceleratorContext
 
 
-if TYPE_CHECKING:
-    import cupy as cp
-
 T = TypeVar("T")
 
 
@@ -100,25 +98,21 @@ def remove_gpu_future(fut_id: int) -> None:
         if fut_id in GPUFuture.gpu_futures:
             GPUFuture.gpu_futures.pop(fut_id).destroy_event()
 
-    def __init__(
-        self, buf: Any, fut_id: int, stream: Optional["cp.cuda.Stream"] = None
-    ):
+    def __init__(self, buf: Any, fut_id: int, stream: Any = None):
         """
         Initialize a GPU future on the given stream.
 
         Args:
             buf: The buffer to return when the future is resolved.
             fut_id: The future ID to cache the future.
-            stream: The CUDA stream to record the event on, this event is waited
+            stream: The torch stream to record the event on, this event is waited
                 on when the future is resolved. If None, the current stream is used.
         """
-        import cupy as cp
-
         if stream is None:
-            stream = cp.cuda.get_current_stream()
+            stream = AcceleratorContext.get().current_stream()
 
         self._buf = buf
-        self._event = cp.cuda.Event()
+        self._event = AcceleratorContext.get().create_event()
         self._event.record(stream)
         self._fut_id = fut_id
         self._waited: bool = False
@@ -131,9 +125,7 @@ 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 = AcceleratorContext.get().current_stream()
         if not self._waited:
             self._waited = True
             current_stream.wait_event(self._event)
@@ -146,11 +138,7 @@ def destroy_event(self) -> None:
         """
         Destroy the CUDA event associated with this future.
         """
-        import cupy as cp
-
         if self._event is None:
             return
 
-        cp.cuda.runtime.eventDestroy(self._event.ptr)
-        self._event.ptr = 0
         self._event = None
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 e31645ba0514..2efd928203c4 100644
--- a/python/ray/dag/tests/experimental/test_torch_tensor_dag.py
+++ b/python/ray/dag/tests/experimental/test_torch_tensor_dag.py
@@ -18,7 +18,6 @@
     Communicator,
     TorchTensorAllocator,
 )
-from ray.experimental.channel.utils import get_devices
 from ray.experimental.channel.torch_tensor_type import TorchTensorType
 from ray.experimental.channel.nccl_group import _NcclGroup
 from ray._private.test_utils import (
@@ -28,6 +27,7 @@
 
 from ray.tests.conftest import *  # noqa
 from ray.experimental.util.types import ReduceOp
+from ray.experimental.channel.accelerator_context import AcceleratorContext
 
 logger = logging.getLogger(__name__)
 
@@ -40,7 +40,7 @@
 @ray.remote
 class TorchTensorWorker:
     def __init__(self):
-        self.device = get_devices()[0]
+        self.device = AcceleratorContext.get().get_default_device()
 
     def init_distributed(self, world_size, rank):
         torch.distributed.init_process_group(
@@ -214,9 +214,8 @@ def test_torch_tensor_as_dag_input(ray_start_regular):
 def test_torch_tensor_nccl(
     ray_start_regular, monkeypatch, enable_profiling, overlap_gpu_communication
 ):
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     monkeypatch.setattr(
         ray.dag.constants, "RAY_CGRAPH_ENABLE_PROFILING", enable_profiling
@@ -316,9 +315,8 @@ def test_torch_tensor_shm(ray_start_regular):
 @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
 @pytest.mark.parametrize("num_gpus", [[0, 0], [1, 0], [0, 1], [1, 1], [0.5, 0.5]])
 def test_torch_tensor_auto(ray_start_regular, num_gpus):
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     sender = TorchTensorWorker.options(num_cpus=0, num_gpus=num_gpus[0]).remote()
     receiver = TorchTensorWorker.options(num_cpus=0, num_gpus=num_gpus[1]).remote()
@@ -376,9 +374,8 @@ def test_torch_tensor_auto(ray_start_regular, num_gpus):
     indirect=["ray_start_regular"],
 )
 def test_torch_tensor_nccl_overlap_timed(ray_start_regular, overlap_gpu_communication):
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) >= 4
-    ), "This test requires at least 4 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 4:
+        pytest.skip("This test requires at least 4 GPUs")
 
     worker_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)
     num_senders = 3
@@ -422,9 +419,8 @@ def test_torch_tensor_nccl_disallows_driver(ray_start_regular):
     and output nodes cannot have a TorchTensorType(transport="nccl")
     annotation.
     """
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)
 
@@ -462,9 +458,8 @@ def test_torch_tensor_nccl_disallows_driver(ray_start_regular):
 @pytest.mark.skipif(not USE_GPU, reason="Skipping GPU Test")
 @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
 def test_torch_tensor_custom_comm(ray_start_regular):
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)
 
@@ -476,8 +471,6 @@ class TestNcclGroup(Communicator):
         A custom NCCL group for testing. This is a simple wrapper around `_NcclGroup`.
         """
 
-        import cupy as cp
-
         def __init__(self, world_size, comm_id, actor_handles):
             self._world_size = world_size
             self._comm_id = comm_id
@@ -490,7 +483,7 @@ def initialize(self, rank: int) -> None:
                 self._comm_id,
                 rank,
                 self._actor_handles,
-                torch.cuda.current_stream().cuda_stream,
+                AcceleratorContext.get().current_stream(),
             )
 
         def get_rank(self, actor: ray.actor.ActorHandle) -> int:
@@ -555,11 +548,11 @@ def reducescatter(
             recv_buf += 1
 
         @property
-        def recv_stream(self) -> Optional["cp.cuda.ExternalStream"]:
+        def recv_stream(self):
             return self._inner.recv_stream
 
         @property
-        def send_stream(self) -> Optional["cp.cuda.ExternalStream"]:
+        def send_stream(self):
             return self._inner.send_stream
 
         def destroy(self) -> None:
@@ -568,6 +561,10 @@ def destroy(self) -> None:
         def get_transport_name(self) -> str:
             return "nccl"
 
+        @classmethod
+        def generate_communicator_id(self) -> str:
+            return self._inner.generate_communicator_id()
+
     from cupy.cuda import nccl
 
     comm_id = nccl.get_unique_id()
@@ -595,9 +592,8 @@ def get_transport_name(self) -> str:
 @pytest.mark.skipif(not USE_GPU, reason="Skipping GPU Test")
 @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
 def test_torch_tensor_custom_comm_inited(ray_start_regular):
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
     runtime_env = {
         "env_vars": {
             "MASTER_ADDR": socket.gethostbyname(socket.gethostname()),
@@ -624,8 +620,6 @@ class InitedNcclGroup(Communicator):
         A custom NCCL group based on existing torch.distributed setup.
         """
 
-        import cupy as cp
-
         def __init__(self, world_size, actor_handles):
             self._world_size = world_size
             self._actor_handles = actor_handles
@@ -637,7 +631,7 @@ def initialize(self, rank: int) -> None:
                 rank == expected_rank
             ), f"NCCL actor's rank {rank} does not match expected rank {expected_rank}"
             self._rank = rank
-            self._device = get_devices()[0]
+            self._device = AcceleratorContext.get().get_default_device()
 
         def get_rank(self, actor: ray.actor.ActorHandle) -> int:
             actor_ids = [a._ray_actor_id for a in self._actor_handles]
@@ -694,16 +688,12 @@ def reducescatter(
             raise NotImplementedError
 
         @property
-        def recv_stream(self) -> Optional["cp.cuda.ExternalStream"]:
-            import cupy as cp
-
-            return cp.cuda.get_current_stream()
+        def recv_stream(self):
+            return AcceleratorContext.get().current_stream()
 
         @property
-        def send_stream(self) -> Optional["cp.cuda.ExternalStream"]:
-            import cupy as cp
-
-            return cp.cuda.get_current_stream()
+        def send_stream(self):
+            return AcceleratorContext.get().current_stream()
 
         def destroy(self) -> None:
             pass
@@ -711,6 +701,10 @@ def destroy(self) -> None:
         def get_transport_name(self) -> str:
             return "nccl"
 
+        @classmethod
+        def generate_communicator_id(self) -> str:
+            return self._inner.generate_communicator_id()
+
     nccl_group = InitedNcclGroup(2, [sender, receiver])
 
     with InputNode() as inp:
@@ -740,9 +734,8 @@ def get_transport_name(self) -> str:
     [["auto", "nccl"], ["custom", "nccl"], ["auto", "nccl"], ["custom", "custom"]],
 )
 def test_torch_tensor_default_comm(ray_start_regular, transports):
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 2
-    ), "This test requires at least 3 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 3:
+        pytest.skip("This test requires at least 3 GPUs")
     runtime_env = {
         "env_vars": {
             "MASTER_ADDR": socket.gethostbyname(socket.gethostname()),
@@ -771,8 +764,6 @@ class InitedNcclGroup(Communicator):
         A custom NCCL group based on existing torch.distributed setup.
         """
 
-        import cupy as cp
-
         def __init__(self, world_size, actor_handles):
             self._world_size = world_size
             self._actor_handles = actor_handles
@@ -784,7 +775,7 @@ def initialize(self, rank: int) -> None:
                 rank == expected_rank
             ), f"NCCL actor's rank {rank} does not match expected rank {expected_rank}"
             self._rank = rank
-            self._device = get_devices()[0]
+            self._device = AcceleratorContext.get().get_default_device()
 
         def get_rank(self, actor: ray.actor.ActorHandle) -> int:
             actor_ids = [a._ray_actor_id for a in self._actor_handles]
@@ -841,16 +832,12 @@ def reducescatter(
             raise NotImplementedError
 
         @property
-        def recv_stream(self) -> Optional["cp.cuda.ExternalStream"]:
-            import cupy as cp
-
-            return cp.cuda.get_current_stream()
+        def recv_stream(self):
+            return AcceleratorContext.get().current_stream()
 
         @property
-        def send_stream(self) -> Optional["cp.cuda.ExternalStream"]:
-            import cupy as cp
-
-            return cp.cuda.get_current_stream()
+        def send_stream(self):
+            return AcceleratorContext.get().current_stream()
 
         def destroy(self) -> None:
             pass
@@ -858,6 +845,10 @@ def destroy(self) -> None:
         def get_transport_name(self) -> str:
             return "nccl"
 
+        @classmethod
+        def generate_communicator_id(self) -> str:
+            return self._inner.generate_communicator_id()
+
     default_comm = InitedNcclGroup(3, [worker0, worker1, worker2])
     custom_comm = InitedNcclGroup(3, [worker0, worker1, worker2])
     custom_comm_count = 0
@@ -902,9 +893,8 @@ def get_transport_name(self) -> str:
 @pytest.mark.skipif(not USE_GPU, reason="Skipping GPU Test")
 @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
 def test_torch_tensor_invalid_custom_comm(ray_start_regular):
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
     runtime_env = {
         "env_vars": {
             "MASTER_ADDR": socket.gethostbyname(socket.gethostname()),
@@ -931,8 +921,6 @@ class UserCreatedNcclGroup(Communicator):
         A custom NCCL group based on existing torch.distributed setup.
         """
 
-        import cupy as cp
-
         def __init__(self, world_size, actor_handles):
             self._world_size = world_size
             self._actor_handles = actor_handles
@@ -944,7 +932,7 @@ def initialize(self, rank: int) -> None:
                 rank == expected_rank
             ), f"NCCL actor's rank {rank} does not match expected rank {expected_rank}"
             self._rank = rank
-            self._device = get_devices()[0]
+            self._device = AcceleratorContext.get().get_default_device()
 
         def get_rank(self, actor: ray.actor.ActorHandle) -> int:
             actor_ids = [a._ray_actor_id for a in self._actor_handles]
@@ -1001,16 +989,12 @@ def reducescatter(
             raise NotImplementedError
 
         @property
-        def recv_stream(self) -> Optional["cp.cuda.ExternalStream"]:
-            import cupy as cp
-
-            return cp.cuda.get_current_stream()
+        def recv_stream(self):
+            return AcceleratorContext.get().current_stream()
 
         @property
-        def send_stream(self) -> Optional["cp.cuda.ExternalStream"]:
-            import cupy as cp
-
-            return cp.cuda.get_current_stream()
+        def send_stream(self):
+            return AcceleratorContext.get().current_stream()
 
         def destroy(self) -> None:
             pass
@@ -1018,6 +1002,10 @@ def destroy(self) -> None:
         def get_transport_name(self) -> str:
             return "nccl"
 
+        @classmethod
+        def generate_communicator_id(self) -> str:
+            return self._inner.generate_communicator_id()
+
     comm2 = UserCreatedNcclGroup(2, [sender, receiver])
     comm1 = UserCreatedNcclGroup(1, [sender])
 
@@ -1047,12 +1035,8 @@ def get_transport_name(self) -> str:
 @pytest.mark.skipif(not USE_GPU, reason="Skipping GPU Test")
 @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
 def test_torch_tensor_nccl_static_shape(ray_start_regular):
-    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"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)
 
@@ -1086,9 +1070,8 @@ def test_torch_tensor_nccl_static_shape(ray_start_regular):
 @pytest.mark.skipif(not USE_GPU, reason="Skipping GPU Test")
 @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
 def test_torch_tensor_nccl_direct_return(ray_start_regular):
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_gpus=1)
 
@@ -1125,9 +1108,8 @@ def test_torch_tensor_nccl_nested_dynamic(ray_start_regular):
     Test nested torch.Tensor passed via NCCL. Its shape and dtype is
     dynamically declared, and there may be multiple tensors.
     """
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_gpus=1)
 
@@ -1161,9 +1143,8 @@ def test_torch_tensor_exceptions(
     """
     Test exceptions being thrown by a NCCL sending task's execution.
     """
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_gpus=1)
 
@@ -1244,9 +1225,8 @@ def test_torch_tensor_exceptions2(
     """
     Test exceptions being thrown by a NCCL sending task's write operation.
     """
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_gpus=1)
     sender = actor_cls.remote()
@@ -1279,12 +1259,45 @@ def test_torch_tensor_exceptions2(
         ref = compiled_dag.execute(2)
 
 
+@pytest.mark.skipif(not USE_GPU, reason="Skipping GPU Test")
+@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 2}], indirect=True)
+def test_torch_tensor_exceptions3(
+    ray_start_regular,
+):
+    """
+    Test exception when creating a communicator group with
+    actors using different accelerators.
+    """
+
+    sender = TorchTensorWorker.options(num_gpus=1).remote()
+    receiver = TorchTensorWorker.options(num_gpus=0).remote()
+
+    with InputNode() as inp:
+        dag = sender.send_int.bind(inp)
+        dag = dag.with_tensor_transport(
+            transport="nccl",
+            _direct_return=True,
+            _static_shape=True,
+        )
+        dag = receiver.recv.bind(dag)
+
+    with pytest.raises(
+        ValueError,
+        match=(
+            r"Actor Actor\(TorchTensorWorker, .*?\) returns a tensor with type hint "
+            r'TorchTensor\(transport="nccl"\) or '
+            r"TorchTensor\(transport=nccl_group_handle\) "
+            r"but actor does not have an accelerator assigned by Ray\."
+        ),
+    ):
+        dag.experimental_compile()
+
+
 @pytest.mark.skipif(not USE_GPU, reason="Skipping GPU Test")
 @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
 def test_torch_tensor_explicit_communicator(ray_start_regular):
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)
 
@@ -1341,9 +1354,8 @@ def test_torch_tensor_nccl_collective_ops(ray_start_regular, operation, reduce_o
     """
     Test basic collective operations.
     """
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)
 
@@ -1441,9 +1453,8 @@ def test_torch_tensor_nccl_all_reduce_get_partial(ray_start_regular):
     """
     Test getting partial results from an all-reduce does not hang.
     """
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)
 
@@ -1484,9 +1495,8 @@ def test_torch_tensor_nccl_all_reduce_wrong_shape(ray_start_regular):
     """
     Test an error is thrown when an all-reduce takes tensors of wrong shapes.
     """
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)
 
@@ -1534,9 +1544,8 @@ def test_torch_tensor_nccl_all_reduce_custom_comm(ray_start_regular):
     """
     Test all-reduce works with a custom communicator.
     """
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)
 
@@ -1550,8 +1559,6 @@ class TestNcclGroup(Communicator):
         A custom NCCL group for testing. This is a simple wrapper around `_NcclGroup`.
         """
 
-        import cupy as cp
-
         def __init__(self, world_size, comm_id, actor_handles):
             self._world_size = world_size
             self._comm_id = comm_id
@@ -1564,7 +1571,7 @@ def initialize(self, rank: int) -> None:
                 self._comm_id,
                 rank,
                 self._actor_handles,
-                torch.cuda.current_stream().cuda_stream,
+                AcceleratorContext.get().current_stream(),
             )
 
         def get_rank(self, actor: ray.actor.ActorHandle) -> int:
@@ -1629,11 +1636,11 @@ def reducescatter(
             recv_buf += 1
 
         @property
-        def recv_stream(self) -> Optional["cp.cuda.ExternalStream"]:
+        def recv_stream(self):
             return self._inner.recv_stream
 
         @property
-        def send_stream(self) -> Optional["cp.cuda.ExternalStream"]:
+        def send_stream(self):
             return self._inner.send_stream
 
         def destroy(self) -> None:
@@ -1642,6 +1649,10 @@ def destroy(self) -> None:
         def get_transport_name(self) -> str:
             return "nccl"
 
+        @classmethod
+        def generate_communicator_id(self) -> str:
+            return self._inner.generate_communicator_id()
+
     comm_id = nccl.get_unique_id()
     nccl_group = TestNcclGroup(2, comm_id, workers)
     with InputNode() as inp:
@@ -1688,9 +1699,8 @@ def test_torch_tensor_nccl_all_reduce_scheduling(ray_start_regular):
     actor 0 starts sending t, then actor 1 waits for actor 0 to join the all-reduce
     while actor 1 waits for actor 0 to receive t.
     """
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)
 
@@ -1730,9 +1740,8 @@ def test_nccl_all_reduce_with_class_method_output_node(ray_start_regular):
     """
     Test all-reduce with class method output node.
     """
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)
 
@@ -1799,9 +1808,8 @@ def recv(self, tensor):
 @pytest.mark.skipif(not USE_GPU, reason="Skipping GPU Test")
 @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
 def test_torch_nccl_channel_with_local_reader(ray_start_regular):
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
 
     actor_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)
 
@@ -1834,9 +1842,8 @@ def test_torch_nccl_channel_with_local_reader(ray_start_regular):
 @pytest.mark.skipif(not USE_GPU, reason="Skipping GPU Test")
 @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
 def test_torch_nccl_channel_with_two_local_readers(ray_start_regular):
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 1
-    ), "This test requires at least 2 GPUs"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 2:
+        pytest.skip("This test requires at least 2 GPUs")
     actor_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)
 
     w1 = actor_cls.remote()
@@ -1869,9 +1876,8 @@ def test_torch_nccl_channel_with_two_local_readers(ray_start_regular):
 @pytest.mark.skipif(not USE_GPU, reason="Skipping GPU Test")
 @pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
 def test_torch_nccl_channel_with_all_local_readers(ray_start_regular):
-    assert (
-        sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) > 0
-    ), "This test requires at least 1 GPU"
+    if sum(node["Resources"].get("GPU", 0) for node in ray.nodes()) < 1:
+        pytest.skip("This test requires at least 1 GPU")
     actor_cls = TorchTensorWorker.options(num_cpus=0, num_gpus=1)
 
     worker = actor_cls.remote()
diff --git a/python/ray/experimental/channel/accelerator_context.py b/python/ray/experimental/channel/accelerator_context.py
new file mode 100644
index 000000000000..d4f656078aa0
--- /dev/null
+++ b/python/ray/experimental/channel/accelerator_context.py
@@ -0,0 +1,201 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type, ContextManager
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+_global_custom_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This includes stream management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, communicator_cls: Type[Communicator]):
+        """
+        Initializes an accelerator context with the specified torch device module
+        and communicator class.
+
+        Args:
+            torch_module_name: Name of the torch device module (e.g., "cuda", "cpu").
+            communicator_cls: Class used to handle communication.
+        """
+
+        # The name of the torch module (e.g., 'cuda', 'npu')
+        self._torch_module_name: str = torch_module_name
+        # The Communicator class used to manage communication
+        self._communicator_cls: Type[Communicator] = communicator_cls
+
+        # Save the torch module corresponding to the device module name.
+        self._torch_mod = importlib.import_module(f"torch.{torch_module_name}")
+
+    @staticmethod
+    def get() -> "AcceleratorContext":
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+
+        global _default_accelerator_context, _global_custom_context
+
+        with _accelerator_context_lock:
+            if _global_custom_context is not None:
+                return _global_custom_context
+
+            if _default_accelerator_context is None:
+                if len(ray.get_gpu_ids()) > 0:
+                    from ray.experimental.channel.nccl_group import _NcclGroup
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cuda", _NcclGroup
+                    )
+                else:
+                    from ray.experimental.channel.cpu_communicator import (
+                        CPUCommunicator,
+                    )
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cpu", CPUCommunicator
+                    )
+
+            return _default_accelerator_context
+
+    @staticmethod
+    def set(accelerator_context: "AcceleratorContext") -> None:
+        """
+        Overwrites the default accelerator context.
+
+        Args:
+            accelerator_context: The context to register.
+        """
+        global _global_custom_context
+
+        # Accelerator context is registered.
+        _global_custom_context = accelerator_context
+
+    def get_default_device(self) -> "torch.device":
+        """
+        Returns the default device used by the compiled graph.
+
+        Currently, the default device for the compiled graph is the first visible
+        device. By default, it returns the device with logical ID 0.
+
+        Returns:
+            torch.device: The default device.
+        """
+        import torch
+
+        if self._torch_module_name == "cpu":
+            return torch.device("cpu")
+
+        return torch.device(f"{self._torch_module_name}:0")
+
+    def get_device_context(self, device: "torch.device") -> ContextManager:
+        """
+        Retrieves the context manager for the specified accelerator device.
+        There is no device context for CPU, returning a nullcontext.
+
+        Args:
+            device: The target device for which the context manager is required.
+
+        Returns:
+            ContextManager: A context manager specific to the device type.
+        """
+        if device.type == "cpu":
+            return nullcontext()
+
+        return self._torch_mod.device(device)
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        return self._torch_mod.current_stream()
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        return self._torch_mod.Event()
+
+    def generate_communicator_id(self) -> str:
+        """
+        Generates a communication identifier for communication group.
+        """
+        return self._communicator_cls.generate_communicator_id()
+
+    def create_communicator(self, *args, **kwargs) -> Communicator:
+        """
+        Creates a communication group for collective operations.
+        """
+        return self._communicator_cls(*args, **kwargs)
+
+    @property
+    def module_name(self) -> str:
+        """
+        Gets the name of the torch module backing the accelerator.
+        """
+        return self._torch_module_name
+
+    @property
+    def communicator_cls(self) -> Optional[Type[Communicator]]:
+        """
+        Returns the communicator class.
+        """
+        return self._communicator_cls
+
+    @property
+    def accelerator_count(self) -> int:
+        """
+        Returns the number of accelerators assigned by ray.
+        """
+        if self._torch_module_name == "cuda":
+            return len(ray.get_gpu_ids())
+        else:
+            accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
+            return len(accelerator_ids.get(self._torch_module_name.upper(), []))
+
+
+def register_accelerator_context(
+    torch_module_name: str, communicator_cls: Type[Communicator]
+):
+    """
+    Registers the accelerator context with the specified device type and communicator.
+
+    Args:
+        torch_module_name: The name of the device module under torch.
+        communicator_cls: The communicator class associated with the device.
+    """
+    accelerator_context = AcceleratorContext(torch_module_name, communicator_cls)
+    AcceleratorContext.set(accelerator_context)
+
+
+def is_accelerator_context_registered():
+    """
+    Checks whether a custom accelerator context has been registered.
+
+    Returns:
+        bool: True if a custom accelerator context is registered
+              (_global_custom_context is not None), False otherwise.
+    """
+    if _global_custom_context is not None:
+        return True
+    return False
diff --git a/python/ray/experimental/channel/common.py b/python/ray/experimental/channel/common.py
index f08a42ed34ca..d6de8cc7a9f4 100644
--- a/python/ray/experimental/channel/common.py
+++ b/python/ray/experimental/channel/common.py
@@ -19,7 +19,8 @@
 import ray
 import ray.exceptions
 from ray.experimental.channel.communicator import Communicator
-from ray.experimental.channel.utils import get_devices
+from ray.experimental.channel.communicator_handle import CommunicatorHandle
+from ray.experimental.channel.accelerator_context import AcceleratorContext
 from ray.experimental.channel.serialization_context import _SerializationContext
 from ray.util.annotations import DeveloperAPI, PublicAPI
 
@@ -127,6 +128,8 @@ class ChannelContext:
     def __init__(self):
         # Used for the torch.Tensor NCCL transport.
         self.communicators: Dict[str, "Communicator"] = {}
+        # Used for driver process to store actors in the communicator.
+        self.communicator_handles: Dict[str, "CommunicatorHandle"] = {}
 
     @staticmethod
     def get_current() -> "ChannelContext":
@@ -163,7 +166,7 @@ def torch_available(self) -> bool:
     @property
     def torch_device(self) -> "torch.device":
         if self._torch_device is None:
-            self._torch_device = get_devices()[0]
+            self._torch_device = AcceleratorContext.get().get_default_device()
 
         return self._torch_device
 
diff --git a/python/ray/experimental/channel/communicator.py b/python/ray/experimental/channel/communicator.py
index 4a8a4fbc3af0..7be8d251acce 100644
--- a/python/ray/experimental/channel/communicator.py
+++ b/python/ray/experimental/channel/communicator.py
@@ -6,7 +6,6 @@
 from ray.util.annotations import DeveloperAPI
 
 if TYPE_CHECKING:
-    import cupy as cp
     import torch
 
 
@@ -108,17 +107,17 @@ def recv(
 
     @property
     @abstractmethod
-    def recv_stream(self) -> Optional["cp.cuda.ExternalStream"]:
+    def recv_stream(self):
         """
-        Return the cuda stream used for receiving tensors.
+        Return the torch stream context used for receiving tensors.
         """
         raise NotImplementedError
 
     @property
     @abstractmethod
-    def send_stream(self) -> Optional["cp.cuda.ExternalStream"]:
+    def send_stream(self):
         """
-        Return the cuda stream used for sending tensors.
+        Return the torch stream context used for sending tensors.
         """
         raise NotImplementedError
 
@@ -190,3 +189,11 @@ def get_transport_name(self) -> str:
         Return the type of the communicator (gpu or cpu).
         """
         raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def generate_communicator_id(cls) -> str:
+        """
+        Return the unique id of the communicator.
+        """
+        raise NotImplementedError
diff --git a/python/ray/experimental/channel/communicator_handle.py b/python/ray/experimental/channel/communicator_handle.py
new file mode 100644
index 000000000000..26dd000ad98c
--- /dev/null
+++ b/python/ray/experimental/channel/communicator_handle.py
@@ -0,0 +1,27 @@
+from typing import List
+import ray
+
+
+class CommunicatorHandle:
+    """
+    A lightweight communicator handle used by the driver to store handles to
+    the actors in the communicator.
+    """
+
+    def __init__(
+        self,
+        actor_handles: List["ray.actor.ActorHandle"],
+    ):
+        """
+        Initializes the CommunicatorHandle with the given actor handles.
+
+        Args:
+            actor_handles: A list of actor handles to be stored.
+        """
+        self._actor_handles = actor_handles
+
+    def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        """
+        Retuan all actor handles in this communicator.
+        """
+        return self._actor_handles
diff --git a/python/ray/experimental/channel/conftest.py b/python/ray/experimental/channel/conftest.py
index 8e6a82e1bb62..5d69ff163b2b 100644
--- a/python/ray/experimental/channel/conftest.py
+++ b/python/ray/experimental/channel/conftest.py
@@ -8,6 +8,7 @@
 import ray
 import ray.dag
 import ray.experimental.channel as ray_channel
+from ray.experimental.channel import nccl_group
 from ray.experimental.channel.communicator import TorchTensorAllocator
 from ray.experimental.util.types import Device
 
@@ -66,8 +67,11 @@ class MockCudaStream:
     def __init__(self):
         self.cuda_stream = 0
 
+    def synchronize(self):
+        pass
 
-class MockNcclGroup(ray_channel.nccl_group._NcclGroup):
+
+class MockNcclGroup(nccl_group._NcclGroup):
     """
     Mock the internal _NcclGroup to use a barrier actor instead of a NCCL group
     for communication.
@@ -133,7 +137,7 @@ def start_nccl_mock():
     cp_patcher.start()
 
     # Mock send/recv ops to use an actor instead of NCCL.
-    ray.experimental.channel.torch_tensor_nccl_channel._NcclGroup = MockNcclGroup
+    ray.experimental.channel.nccl_group._NcclGroup = MockNcclGroup
 
     # PyTorch mocks.
     stream_patcher = mock.patch(
diff --git a/python/ray/experimental/channel/cpu_communicator.py b/python/ray/experimental/channel/cpu_communicator.py
index 8b7c7665ed1b..0ee1679dad68 100644
--- a/python/ray/experimental/channel/cpu_communicator.py
+++ b/python/ray/experimental/channel/cpu_communicator.py
@@ -199,3 +199,9 @@ def recv_stream(self):
 
     def send_stream(self):
         raise NotImplementedError
+
+    @classmethod
+    def generate_communicator_id(cls) -> str:
+        import uuid
+
+        return str(uuid.uuid4())
diff --git a/python/ray/experimental/channel/nccl_group.py b/python/ray/experimental/channel/nccl_group.py
index c3d1798a5674..520df23b0cc0 100644
--- a/python/ray/experimental/channel/nccl_group.py
+++ b/python/ray/experimental/channel/nccl_group.py
@@ -6,10 +6,9 @@
 from ray.exceptions import RayChannelError
 from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
 from ray.experimental.util.types import ReduceOp
-from ray.experimental.channel.utils import get_devices
+from ray.experimental.channel.accelerator_context import AcceleratorContext
 
 if TYPE_CHECKING:
-    import cupy as cp
     import torch
 
 
@@ -33,7 +32,7 @@ def __init__(
         comm_id: tuple,
         rank: Optional[int],
         actor_handles: List["ray.actor.ActorHandle"],
-        cuda_stream: Optional[int],
+        cuda_stream: Optional["torch.cuda.Stream"],
         use_communication_streams: bool = False,
     ):
         """
@@ -93,29 +92,21 @@ def __init__(
             # Driver does not have a rank.
             self._comm = None
 
-        self._cuda_stream: Optional["cp.cuda.ExternalStream"] = None
-        self._send_stream: Optional["cp.cuda.ExternalStream"] = None
-        self._recv_stream: Optional["cp.cuda.ExternalStream"] = None
+        self._cuda_stream: Optional["torch.cuda.Stream"] = None
+        self._send_stream: Optional["torch.cuda.Stream"] = None
+        self._recv_stream: Optional["torch.cuda.Stream"] = None
         if cuda_stream is not None:
             assert rank is not None, "NCCL actor has no rank assigned"
-
-            import cupy as cp
-
-            # TODO(swang): Allow default device to be overridden.
-            device = get_devices()[0]
-            self._cuda_stream = cp.cuda.ExternalStream(
-                cuda_stream, device_id=device.index
-            )
+            self._cuda_stream = cuda_stream
 
             if use_communication_streams:
                 import torch
 
-                self._send_stream = cp.cuda.ExternalStream(
-                    torch.cuda.Stream().cuda_stream, device_id=device.index
-                )
-                self._recv_stream = cp.cuda.ExternalStream(
-                    torch.cuda.Stream().cuda_stream, device_id=device.index
-                )
+                # TODO(swang): Allow default device to be overridden.
+                device = AcceleratorContext.get().get_default_device()
+
+                self._send_stream = torch.cuda.Stream(device=device)
+                self._recv_stream = torch.cuda.Stream(device=device)
             else:
                 self._send_stream = self._cuda_stream
                 self._recv_stream = self._cuda_stream
@@ -189,7 +180,7 @@ def send(self, buf: "torch.Tensor", peer_rank: int) -> None:
             buf.numel(),
             self.nccl_util.get_nccl_tensor_dtype(buf),
             peer_rank,
-            self._send_stream.ptr,
+            self._send_stream.cuda_stream,
         )
 
     def recv(
@@ -228,7 +219,7 @@ def recv(
                 buf.numel(),
                 self.nccl_util.get_nccl_tensor_dtype(buf),
                 peer_rank,
-                self._recv_stream.ptr,
+                self._recv_stream.cuda_stream,
             )
         else:
             self._comm.recv(
@@ -236,7 +227,7 @@ def recv(
                 buf.numel(),
                 self.nccl_util.get_nccl_tensor_dtype(buf),
                 peer_rank,
-                self._recv_stream.ptr,
+                self._recv_stream.cuda_stream,
             )
 
             # Buffer values are undefined if NCCL ops are aborted. Therefore, we
@@ -290,7 +281,7 @@ def allgather(
             self.nccl_util.get_tensor_ptr(recv_buf),
             send_buf.numel(),
             self.nccl_util.get_nccl_tensor_dtype(send_buf),
-            self._cuda_stream.ptr,
+            self._cuda_stream.cuda_stream,
         ]
         self._exec_collective(
             send_buf,
@@ -311,7 +302,7 @@ def allreduce(
             send_buf.numel(),
             self.nccl_util.get_nccl_tensor_dtype(send_buf),
             op.value,
-            self._cuda_stream.ptr,
+            self._cuda_stream.cuda_stream,
         ]
         self._exec_collective(
             send_buf,
@@ -332,7 +323,7 @@ def reducescatter(
             recv_buf.numel(),
             self.nccl_util.get_nccl_tensor_dtype(send_buf),
             op.value,
-            self._cuda_stream.ptr,
+            self._cuda_stream.cuda_stream,
         ]
         self._exec_collective(
             send_buf,
@@ -342,12 +333,16 @@ def reducescatter(
         )
 
     @property
-    def recv_stream(self) -> Optional["cp.cuda.ExternalStream"]:
-        return self._recv_stream
+    def recv_stream(self):
+        import torch
+
+        return torch.cuda.StreamContext(self._recv_stream)
 
     @property
-    def send_stream(self) -> Optional["cp.cuda.ExternalStream"]:
-        return self._send_stream
+    def send_stream(self):
+        import torch
+
+        return torch.cuda.StreamContext(self._send_stream)
 
     def destroy(self) -> None:
         """
@@ -371,3 +366,9 @@ def destroy(self) -> None:
 
     def get_transport_name(self) -> str:
         return "nccl"
+
+    @classmethod
+    def generate_communicator_id(cls) -> str:
+        from cupy.cuda import nccl
+
+        return nccl.get_unique_id()
diff --git a/python/ray/experimental/channel/serialization_context.py b/python/ray/experimental/channel/serialization_context.py
index 3a0deeab8eec..947181a04842 100644
--- a/python/ray/experimental/channel/serialization_context.py
+++ b/python/ray/experimental/channel/serialization_context.py
@@ -2,6 +2,7 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Set, Tuple, Union
 
 from ray.experimental.util.types import Device
+from ray.experimental.channel.accelerator_context import AcceleratorContext
 
 if TYPE_CHECKING:
     import numpy as np
@@ -120,7 +121,7 @@ def serialize_to_numpy_or_scalar(
         # CPU and another from CPU to shared memory. Ideally we should elide
         # the first copy and memcpy directly from GPU to the shared memory
         # buffer.
-        if tensor_device_type == "cuda":
+        if tensor_device_type != "cpu":
             tensor = tensor.to("cpu")
 
         # Numpy does not have an equivalent dtype for all torch dtypes, so
@@ -169,12 +170,12 @@ def deserialize_from_numpy_or_scalar(
         if target_device == Device.DEFAULT:
             target_device_type = tensor_device_type
         elif target_device in [Device.GPU, Device.CUDA]:
-            target_device_type = "cuda"
+            target_device_type = AcceleratorContext.get().module_name
         else:
-            target_device_type = "cpu"
+            target_device_type = target_device.value
 
         # TODO(swang): Support local P2P transfers if available.
-        if target_device_type == "cuda":
+        if target_device_type != "cpu":
 
             def convert_numpy_to_tensor(np_array):
                 if not isinstance(np_array, np.ndarray):
diff --git a/python/ray/experimental/channel/torch_tensor_nccl_channel.py b/python/ray/experimental/channel/torch_tensor_nccl_channel.py
index 22707919da51..7951a09c92fe 100644
--- a/python/ray/experimental/channel/torch_tensor_nccl_channel.py
+++ b/python/ray/experimental/channel/torch_tensor_nccl_channel.py
@@ -3,7 +3,7 @@
 import uuid
 from dataclasses import dataclass
 from types import ModuleType
-from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union, Type
 
 import ray
 import ray.util.serialization
@@ -12,10 +12,15 @@
 from ray.experimental.channel.communicator import Communicator
 from ray.experimental.channel.cpu_communicator import CPUCommunicator
 from ray.experimental.channel.intra_process_channel import IntraProcessChannel
-from ray.experimental.channel.nccl_group import _NcclGroup
+from ray.experimental.channel.communicator_handle import CommunicatorHandle
 from ray.experimental.channel.shared_memory_channel import SharedMemoryType
 from ray.experimental.channel.torch_tensor_type import TorchTensorType
 from ray.util.annotations import DeveloperAPI
+from ray.experimental.channel.accelerator_context import (
+    AcceleratorContext,
+    register_accelerator_context,
+    is_accelerator_context_registered,
+)
 
 if TYPE_CHECKING:
     import torch
@@ -399,32 +404,38 @@ def __init__(
         ), f"NCCL group ID ({typ.communicator_id}) must be a str."
         self._typ = typ
 
-        assert self._typ.communicator_id is not None, "No NCCL group specified."
-        self._nccl_group_id: str = self._typ.communicator_id
-        self._nccl_group: "Communicator" = ctx.communicators[self._typ.communicator_id]
-        assert (
-            self._nccl_group is not None
-        ), "ChannelContext.nccl_group is not initialized."
-
         self._static_shape = typ.static_shape
 
-        self._writer_rank = self._nccl_group.get_rank(self._writer)
-        self._reader_ranks = [
-            self._nccl_group.get_rank(reader)
-            for reader, _ in self._reader_and_node_list
-        ]
-
-        if (
-            self._writer_rank is not None
-            and self._writer_rank == self._nccl_group.get_self_rank()
-        ):
-            self._writer_registered = True
+        assert self._typ.communicator_id is not None, "No NCCL group specified."
+        self._nccl_group_id: str = self._typ.communicator_id
 
-        if (
-            self._reader_ranks
-            and self._nccl_group.get_self_rank() in self._reader_ranks
-        ):
-            self._reader_registered = True
+        # If the communicators does not contain the group_id, it means the current
+        # process is the driver, and there’s no need to fetch the nccl_group.
+        if self._typ.communicator_id in ctx.communicators:
+            self._nccl_group: "Communicator" = ctx.communicators[
+                self._typ.communicator_id
+            ]
+            assert (
+                self._nccl_group is not None
+            ), "ChannelContext.nccl_group is not initialized."
+
+            self._writer_rank = self._nccl_group.get_rank(self._writer)
+            self._reader_ranks = [
+                self._nccl_group.get_rank(reader)
+                for reader, _ in self._reader_and_node_list
+            ]
+
+            if (
+                self._writer_rank is not None
+                and self._writer_rank == self._nccl_group.get_self_rank()
+            ):
+                self._writer_registered = True
+
+            if (
+                self._reader_ranks
+                and self._nccl_group.get_self_rank() in self._reader_ranks
+            ):
+                self._reader_registered = True
 
         # If the channel type specifies that the tensor shape is static, then the
         # receiver can allocate buffers without needing to coordinate with the
@@ -448,13 +459,13 @@ def ensure_registered_as_writer(self):
         assert self._nccl_group is not None, "Actor is not part of a NCCL group"
         assert self._writer_registered
         ctx = ChannelContext.get_current()
-        assert ctx.torch_device.type == "cuda"
+        assert ctx.torch_device.type != "cpu"
 
     def ensure_registered_as_reader(self) -> bool:
         assert self._nccl_group is not None, "Actor is not part of a NCCL group"
         assert self._reader_registered
         ctx = ChannelContext.get_current()
-        assert ctx.torch_device.type == "cuda"
+        assert ctx.torch_device.type != "cpu"
 
     def __reduce__(self):
         return (
@@ -642,12 +653,10 @@ def _do_init_communicator(
     use_communication_streams,
     custom_communicator: Optional[Communicator] = None,
 ):
-    import torch
-
     if not custom_communicator:
         assert (
-            ray.get_gpu_ids()
-        ), "Actors participating in NCCL group must have at least one GPU assigned"
+            AcceleratorContext.get().accelerator_count > 0
+        ), "Actors participating in Communication group must have at least one Accelerator assigned"
 
     ctx = ChannelContext.get_current()
     if custom_communicator is not None:
@@ -655,12 +664,12 @@ def _do_init_communicator(
         ctx.communicators[group_id] = custom_communicator
     else:
         # default to NcclGroup
-        ctx.communicators[group_id] = _NcclGroup(
+        ctx.communicators[group_id] = AcceleratorContext.get().create_communicator(
             world_size,
             comm_id,
             rank,
             actor_handles,
-            torch.cuda.current_stream().cuda_stream,
+            AcceleratorContext.get().current_stream(),
             use_communication_streams,
         )
 
@@ -675,14 +684,16 @@ def _do_destroy_communicator(self, group_id):
     # task loop running.
 
 
-def _do_check_has_gpu(self) -> bool:
-    return bool(ray.get_gpu_ids())
+def _do_check_has_accelerators(self) -> str:
+    return AcceleratorContext.get().accelerator_count > 0
+
 
+def do_register_accelerator_context(self, name: str, communicator: Type[Communicator]):
+    register_accelerator_context(name, communicator)
 
-def _do_get_unique_nccl_id(self) -> tuple:
-    from cupy.cuda import nccl
 
-    return nccl.get_unique_id()
+def _do_get_unique_communication_id(self) -> bool:
+    return AcceleratorContext.get().generate_communicator_id()
 
 
 def _get_ranks(
@@ -722,6 +733,8 @@ def _init_communicator(
     actors: List[ray.actor.ActorHandle],
     custom_communicator: Optional[Communicator] = None,
     use_communication_streams: bool = False,
+    accelerator_module_name: Optional[str] = None,
+    accelerator_communicator_cls: Optional[Type[Communicator]] = None,
 ) -> str:
     """
     Initialize a NCCL group with the given actors. If a custom NCCL group is
@@ -733,6 +746,8 @@ def _init_communicator(
         use_communication_streams: Whether to use dedicated send and recv
                 streams for communication. If True, communication and computation
                 can be overlapped to improve performance.
+        accelerator_module_name: Optional name of the accelerator module to use.
+        accelerator_communicator_cls: Optional communicator class for the accelerator.
     """
     ctx = ChannelContext.get_current()
 
@@ -740,16 +755,30 @@ def _init_communicator(
         custom_communicator, CPUCommunicator
     )
 
-    has_gpus = ray.get(
-        [actor.__ray_call__.remote(_do_check_has_gpu) for actor in actors]
+    # Register accelerator context for all actors if accelerator is not default
+    if accelerator_module_name and accelerator_communicator_cls:
+        if is_accelerator_context_registered():
+            ray.get(
+                [
+                    actor.__ray_call__.remote(
+                        do_register_accelerator_context,
+                        accelerator_module_name,
+                        accelerator_communicator_cls,
+                    )
+                    for actor in actors
+                ]
+            )
+
+    has_accelerators = ray.get(
+        [actor.__ray_call__.remote(_do_check_has_accelerators) for actor in actors]
     )
-    for has_gpu, actor in zip(has_gpus, actors):
-        if not has_gpu and not is_cpu_communicator:
+    for has_accelerator, actor in zip(has_accelerators, actors):
+        if not has_accelerator and not is_cpu_communicator:
             raise ValueError(
                 f"Actor {actor} returns a tensor with type hint "
                 'TorchTensor(transport="nccl") or '
-                "TorchTensor(transport=nccl_group_handle)"
-                "but actor does not have a GPU assigned by Ray."
+                "TorchTensor(transport=nccl_group_handle) "
+                "but actor does not have an accelerator assigned by Ray."
             )
 
     actor_ids = {actor._ray_actor_id for actor in actors}
@@ -758,11 +787,8 @@ def _init_communicator(
     # Allocate a communicator ID on one of the actors that will participate in
     # the group. This is in case the driver is not on the same node as one of
     # the NCCL actors.
-    nccl_comm_id = (
-        ray.get(actors[0].__ray_call__.remote(_do_get_unique_nccl_id))
-        if not is_cpu_communicator
-        else str(uuid.uuid4())
-    )
+    comm_id = ray.get(actors[0].__ray_call__.remote(_do_get_unique_communication_id))
+
     # Used to uniquely identify this NCCL group.
     group_id = str(uuid.uuid4())
 
@@ -778,7 +804,7 @@ def _init_communicator(
             _do_init_communicator,
             group_id,
             world_size,
-            nccl_comm_id,
+            comm_id,
             rank,
             actors,
             use_communication_streams,
@@ -797,15 +823,14 @@ def _init_communicator(
     logger.info("NCCL group initialized.")
 
     if custom_communicator is not None:
-        ctx.communicators[group_id] = custom_communicator
+        ctx.communicator_handles[group_id] = CommunicatorHandle(
+            actor_handles=custom_communicator.get_actor_handles(),
+        )
     else:
-        ctx.communicators[group_id] = _NcclGroup(
-            world_size,
-            nccl_comm_id,
-            rank=None,
+        ctx.communicator_handles[group_id] = CommunicatorHandle(
             actor_handles=actors,
-            cuda_stream=None,
         )
+
     return group_id
 
 
@@ -814,10 +839,10 @@ def _destroy_communicator(group_id: str) -> None:
     Destroy the NCCL group with the given ID.
     """
     ctx = ChannelContext.get_current()
-    if group_id not in ctx.communicators:
+    if group_id not in ctx.communicator_handles:
         return
 
-    group = ctx.communicators[group_id]
+    group = ctx.communicator_handles[group_id]
     actors = group.get_actor_handles()
     destroy_tasks = [
         actor.__ray_call__.remote(
@@ -834,4 +859,4 @@ def _destroy_communicator(group_id: str) -> None:
             "may be hung."
         )
 
-    del ctx.communicators[group_id]
+    del ctx.communicator_handles[group_id]
diff --git a/python/ray/experimental/channel/utils.py b/python/ray/experimental/channel/utils.py
index 8f4e5edb2eec..6df57828130a 100644
--- a/python/ray/experimental/channel/utils.py
+++ b/python/ray/experimental/channel/utils.py
@@ -1,10 +1,6 @@
-from typing import TYPE_CHECKING, List, Optional, Tuple
+from typing import List, Optional, Tuple
 
 import ray
-import os
-
-if TYPE_CHECKING:
-    import torch
 
 
 def get_self_actor() -> Optional["ray.actor.ActorHandle"]:
@@ -94,68 +90,3 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
                 lambda self: ray.get_runtime_context().get_node_id()
             )
         )
-
-
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
-    """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
-
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
-
-    device_ids = []
-
-    if len(gpu_ids) > 0:
-        cuda_visible_str = os.environ.get("CUDA_VISIBLE_DEVICES", "")
-        if cuda_visible_str and cuda_visible_str != "NoDevFiles":
-            cuda_visible_list = cuda_visible_str.split(",")
-        else:
-            cuda_visible_list = []
-
-        # By default, there should only be one GPU ID if `use_gpu=True`.
-        # If there are multiple GPUs, return a list of devices.
-        # If using fractional GPUs, these IDs are not guaranteed
-        # to be unique across different processes.
-        for gpu_id in gpu_ids:
-            try:
-                device_ids.append(cuda_visible_list.index(gpu_id))
-            except IndexError:
-                raise RuntimeError(
-                    "CUDA_VISIBLE_DEVICES set incorrectly. "
-                    f"Got {cuda_visible_str}, expected to include {gpu_id}. "
-                    "Did you override the `CUDA_VISIBLE_DEVICES` environment"
-                    " variable? If not, please help file an issue on Github."
-                )
-
-    else:
-        # If called on the driver or outside of Ray Train, return the
-        # 0th device.
-        device_ids.append(0)
-
-    return [torch.device(f"cuda:{device_id}") for device_id in device_ids]
-
-
-def get_devices() -> List["torch.device"]:
-    """Gets the correct torch device list configured for this process.
-
-    Returns a list of torch devices allocated for the current worker.
-    If no devices are assigned, then it returns a list with a single CPU device.
-    """
-
-    import torch
-
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
-    if len(gpu_ids) > 0:
-        return get_cuda_devices()
-    else:
-        return [torch.device("cpu")]
diff --git a/python/ray/experimental/collective/conftest.py b/python/ray/experimental/collective/conftest.py
index 5660ad810f24..464524e1487f 100644
--- a/python/ray/experimental/collective/conftest.py
+++ b/python/ray/experimental/collective/conftest.py
@@ -1,5 +1,5 @@
 import uuid
-from typing import Dict, FrozenSet, List, Optional, Set, Tuple
+from typing import Dict, FrozenSet, List, Optional, Set, Tuple, Type
 
 import torch
 
@@ -17,8 +17,6 @@ class AbstractNcclGroup(Communicator):
     A dummy NCCL group for testing.
     """
 
-    import cupy as cp
-
     def __init__(self, actor_handles: List[ray.actor.ActorHandle]):
         self._actor_handles = actor_handles
         self._rank = None
@@ -74,11 +72,11 @@ def reducescatter(
         raise NotImplementedError
 
     @property
-    def recv_stream(self) -> Optional["cp.cuda.ExternalStream"]:
+    def recv_stream(self):
         return None
 
     @property
-    def send_stream(self) -> Optional["cp.cuda.ExternalStream"]:
+    def send_stream(self):
         return None
 
     def destroy(self) -> None:
@@ -87,6 +85,10 @@ def destroy(self) -> None:
     def get_transport_name(self) -> str:
         return "nccl"
 
+    @classmethod
+    def generate_communicator_id(cls) -> str:
+        pass
+
 
 class MockNcclGroupSet:
     def __init__(self):
@@ -101,6 +103,8 @@ def __call__(
         actors: List["ray.actor.ActorHandle"],
         custom_nccl_group: Optional[Communicator] = None,
         use_communication_streams: bool = False,
+        accelerator_module_name: Optional[str] = None,
+        accelerator_communicator_cls: Optional[Type[Communicator]] = None,
     ) -> str:
         group_id = str(uuid.uuid4())
         self.ids_to_actors_and_custom_comms[group_id] = (
diff --git a/python/ray/util/collective/util.py b/python/ray/util/collective/util.py
index 7fa384901ac5..09f9a8e8045f 100644
--- a/python/ray/util/collective/util.py
+++ b/python/ray/util/collective/util.py
@@ -26,7 +26,7 @@ def set_id(self, uid):
         Initialize the NCCL unique ID for this store.
 
         Args:
-            uid: the unique ID generated via the NCCL get_unique_id API.
+            uid: the unique ID generated via the NCCL generate_communicator_id API.
 
         Returns:
             None


Reviews:
user: login: noemotiovon
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-04T06:47:31Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-04T11:12:48Z
user: login: ruisearch42
body: Trying to better understand the PR
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-06T02:09:26Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-06T02:32:53Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-06T02:37:31Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-06T02:39:41Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-06T02:40:41Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-06T02:43:03Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-06T07:29:49Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-06T08:03:58Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T02:02:40Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T02:03:01Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T02:04:01Z
user: login: ruisearch42
body: Still reviewing ..
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T23:54:52Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T01:06:29Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T01:06:34Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T01:06:56Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T01:08:33Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-10T01:10:41Z
user: login: ruisearch42
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-11T20:42:06Z
user: login: noemotiovon
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T01:56:16Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T12:28:18Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T12:30:37Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T12:30:51Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T12:30:57Z
user: login: ruisearch42
body: Thanks for the change and continued effort.

I will try but not sure will have a detailed review this week due to oncall tasks. Maybe your colleagues could help review/approve first? especially for the NPU part.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T05:35:27Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T07:00:41Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T07:01:00Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T07:01:15Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T07:01:42Z
user: login: noemotiovon
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T03:47:01Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T12:00:04Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T06:37:03Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T06:37:38Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T08:08:57Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T08:09:26Z
user: login: noemotiovon
body: The overall architecture LGTM. Here are some minor suggestions.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T02:23:04Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-28T01:31:22Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-28T01:31:31Z
user: login: ruisearch42
body: Thanks for the PR. Please address Stephanie's comment as well.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-31T22:47:40Z
user: login: stephanie-wang
body: Thanks for this contribution! At a high level, I think it makes a lot of sense to try to abstract out the AcceleratorRuntime.

However, I believe need some iteration on the API calls. It would be useful if you can write a short doc that summarizes all of the APIs that need to be specialized to a specific device and shows how the rest of the code could call them. I think we'll be able to iterate faster on such a doc instead of the code. Happy to discuss synchronously too.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-06T22:16:22Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T08:28:01Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T08:32:19Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T08:37:38Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T08:42:56Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T08:47:57Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T08:52:22Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T08:53:39Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T08:56:02Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T02:24:59Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T02:25:06Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T02:25:34Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T02:25:53Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T02:25:59Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T02:26:07Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T02:26:17Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T02:26:23Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T02:26:38Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T02:26:49Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T02:26:58Z
user: login: noemotiovon
body: LGTM! 
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T07:00:55Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T07:46:26Z
user: login: stephanie-wang
body: Thanks for the contribution! I like the new APIs. Let's check if we can reduce the driver API; eventually it would be good to use this as a single-controller API for Ray GPU objects too.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T01:09:15Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T06:32:44Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T07:17:47Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T07:36:28Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T07:38:59Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T07:43:46Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T08:56:35Z
user: login: stephanie-wang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T17:54:43Z
user: login: stephanie-wang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T17:55:18Z
user: login: stephanie-wang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T17:57:27Z
user: login: stephanie-wang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T17:59:16Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T12:25:16Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T06:20:39Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T06:20:46Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T06:20:54Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T06:21:02Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T06:21:07Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T06:21:25Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T06:21:31Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T06:22:38Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T06:22:50Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T06:23:18Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T06:23:34Z
user: login: stephanie-wang
body: Thanks, this looks good! Please address the last few comments and then we can merge it.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T18:23:09Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T02:38:51Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T02:38:56Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T02:38:59Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T02:39:02Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T02:39:06Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T02:39:09Z
user: login: hipudding
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T02:39:39Z
user: login: noemotiovon
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T03:12:57Z
user: login: stephanie-wang
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T21:42:51Z

Review comments:
diff_hunk: @@ -108,19 +107,19 @@ def recv(
 
     @property
     @abstractmethod
-    def recv_stream(self) -> Optional["cp.cuda.ExternalStream"]:
+    def recv_stream(self):
         """
-        Return the cuda stream used for receiving tensors.
+        Return the torch stream used for receiving tensors.
         """
-        raise NotImplementedError
+        return None
path: python/ray/experimental/channel/communicator.py
user: login: noemotiovon
body: If the method returns None, is it still necessary to mark it as an abstract method?
created_at: 2025-03-04T06:30:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -108,19 +107,19 @@ def recv(
 
     @property
     @abstractmethod
-    def recv_stream(self) -> Optional["cp.cuda.ExternalStream"]:
+    def recv_stream(self):
         """
-        Return the cuda stream used for receiving tensors.
+        Return the torch stream used for receiving tensors.
         """
-        raise NotImplementedError
+        return None
path: python/ray/experimental/channel/communicator.py
user: login: hipudding
body: Keep this class abstrct.
created_at: 2025-03-04T11:12:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -766,14 +767,12 @@ def exec_operation(
         """
         if op_type == _DAGNodeOperationType.READ:
             with _device_context_manager():
-                with self._recv_stream:
-                    return self._read(overlap_gpu_communication)
+                return self._read(overlap_gpu_communication)
         elif op_type == _DAGNodeOperationType.COMPUTE:
             return self._compute(overlap_gpu_communication, class_handle)
         elif op_type == _DAGNodeOperationType.WRITE:
             with _device_context_manager():
-                with self._send_stream:
-                    return self._write()
path: python/ray/dag/compiled_dag_node.py
user: login: ruisearch42
body: Why don't we want to use context manager?
created_at: 2025-03-06T01:52:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -624,12 +623,12 @@ def wrap_and_set_intermediate_future(
         assert self._intermediate_future is None
 
         if wrap_in_gpu_future:
-            future = GPUFuture(val)
+            future = GPUFuture(val, stream)
         else:
             future = ResolvedFuture(val)
         self._intermediate_future = future
 
-    def reset_and_wait_intermediate_future(self) -> Any:
+    def reset_and_wait_intermediate_future(self, stream: Any = None) -> Any:
path: python/ray/dag/compiled_dag_node.py
user: login: ruisearch42
body: update doc for arg
created_at: 2025-03-06T01:53:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -600,15 +596,18 @@ def prepare(self, overlap_gpu_communication: bool = False):
                         nccl_group_id
                     )
                     assert nccl_group is not None
-                    if not isinstance(self._recv_stream, nullcontext):
+                    if self._recv_stream is not None:
                         assert self._recv_stream == nccl_group.recv_stream, (
                             "Currently all torch tensor input channels of a "
                             "Compiled Graph task should use the same recv cuda stream."
                         )
                     self._recv_stream = nccl_group.recv_stream
 
     def wrap_and_set_intermediate_future(
-        self, val: Any, wrap_in_gpu_future: bool
+        self,
+        val: Any,
+        wrap_in_gpu_future: bool,
+        stream: Any = None,
path: python/ray/dag/compiled_dag_node.py
user: login: ruisearch42
body: update doc for stream
created_at: 2025-03-06T01:54:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,81 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class _DriverGroupHolder(Communicator):
+    """
+    Communicator place holder for Driver, Since driver may has no
+    Accelerators, TorchDeviceManager cannot be used.
+    """
path: python/ray/experimental/channel/driver_communicator_holder.py
user: login: ruisearch42
body: Can you elaborate a bit why this is needed?
created_at: 2025-03-06T01:57:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,242 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+class _HcclGroup(Communicator):
+    """
+    Represents an actor's HCCL communicator. This is the default HCCL communicator
+    to be used in Compiled Graph if a custom communicator is not provided.
+
+    This class is not thread-safe.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        comm_id: int,
+        rank: Optional[int],
+        actor_handles: List["ray.actor.ActorHandle"],
+        acl_stream: Optional[Any],
path: python/ray/experimental/channel/hccl_group.py
user: login: ruisearch42
body: What is acl_tream? doc for args
created_at: 2025-03-06T02:04:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,242 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+class _HcclGroup(Communicator):
+    """
+    Represents an actor's HCCL communicator. This is the default HCCL communicator
+    to be used in Compiled Graph if a custom communicator is not provided.
+
+    This class is not thread-safe.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        comm_id: int,
+        rank: Optional[int],
+        actor_handles: List["ray.actor.ActorHandle"],
+        acl_stream: Optional[Any],
+        use_communication_streams: bool = False,
+    ):
+        self._world_size = world_size
+        self._rank: Optional[int] = rank
+        self.hccl: Optional[ModuleType] = None
+        self._actor_handles = actor_handles
+        self._use_communication_streams = use_communication_streams
+
+        if rank is not None:
+            assert "NPU" in ray.cluster_resources(), "HCCL actor has no NPUs assigned"
+            assert acl_stream is not None, "HCCL actor must specify aclrtStream"
+
+            expected_rank = self.get_rank(ray.get_runtime_context().current_actor)
+            assert (
+                rank == expected_rank
+            ), f"HCCL actor's rank {rank} does not match expected rank {expected_rank}"
+
+            from ray.experimental.channel import hccl
+
+            self.hccl = hccl
+
+            self._comm = self.hccl.HCCLCommunicator(world_size, comm_id, rank)
+        else:
+            # Driver does not have a rank.
+            self._comm = None
+
+        self._acl_stream: None
+        self._send_stream: None
+        self._recv_stream: None
+        if acl_stream is not None:
+            assert rank is not None, "HCCL actor has no rank assigned"
+
+            self._acl_stream = acl_stream
+
+            if use_communication_streams:
+                import torch
+                import torch_npu  # noqa: F401
+                from ray.air._internal import torch_utils
+
+                # TODO(swang): Allow default device to be overridden.
+                device = torch_utils.get_devices()[0]
+
+                self._send_stream = torch.npu.Stream(device=device)
+                self._recv_stream = torch.npu.Stream(device=device)
+            else:
+                self._send_stream = self._acl_stream
+                self._recv_stream = self._acl_stream
+
+        self._closed = False
+
+    def initialize(self, rank: int) -> None:
+        pass
+
+    def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        return self._actor_handles
+
+    def get_rank(self, actor: ray.actor.ActorHandle) -> int:
+        actor_ids = [a._ray_actor_id for a in self._actor_handles]
+        try:
+            rank = actor_ids.index(actor._ray_actor_id)
+        except ValueError:
+            raise ValueError("Actor is not in the HCCL group.")
+        return rank
+
+    def get_self_rank(self) -> Optional[int]:
+        return self._rank
+
+    def get_world_size(self) -> int:
+        return self._world_size
+
+    def send(self, buf: "torch.Tensor", peer_rank: int) -> None:
path: python/ray/experimental/channel/hccl_group.py
user: login: ruisearch42
body: many methods in this class look like code duplication from `_NcclGroup` with minor changes. Should they be unified?
created_at: 2025-03-06T02:06:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -766,14 +767,12 @@ def exec_operation(
         """
         if op_type == _DAGNodeOperationType.READ:
             with _device_context_manager():
-                with self._recv_stream:
-                    return self._read(overlap_gpu_communication)
+                return self._read(overlap_gpu_communication)
         elif op_type == _DAGNodeOperationType.COMPUTE:
             return self._compute(overlap_gpu_communication, class_handle)
         elif op_type == _DAGNodeOperationType.WRITE:
             with _device_context_manager():
-                with self._send_stream:
-                    return self._write()
path: python/ray/dag/compiled_dag_node.py
user: login: hipudding
body: This stream context manager is provided by cupy, it only support CUDA devices, So I just remote context manager and specified stream in read and write. 
But I think I can implement a stream context manager in Ray, so cupy is not needed. Which is better?
created_at: 2025-03-06T02:32:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,81 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class _DriverGroupHolder(Communicator):
+    """
+    Communicator place holder for Driver, Since driver may has no
+    Accelerators, TorchDeviceManager cannot be used.
+    """
path: python/ray/experimental/channel/driver_communicator_holder.py
user: login: hipudding
body: I replace all torch.cuda.xxx by TorchDeviceManager, including create a _NcclGroup. But driver will keep a communicator instance itself. Driver will not join the group, but need record some information like world size and actors.
But Driver may has no GPU ( or other accelerators), it will choose the wrong TorchDeviceManager (CPUTorchDeviceManager), and get_communicator will return None. 
Since Driver will not join the group, it use Communicator only for store some information. So I create a Hardware-independent class for driver.
created_at: 2025-03-06T02:37:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,242 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+class _HcclGroup(Communicator):
+    """
+    Represents an actor's HCCL communicator. This is the default HCCL communicator
+    to be used in Compiled Graph if a custom communicator is not provided.
+
+    This class is not thread-safe.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        comm_id: int,
+        rank: Optional[int],
+        actor_handles: List["ray.actor.ActorHandle"],
+        acl_stream: Optional[Any],
path: python/ray/experimental/channel/hccl_group.py
user: login: hipudding
body: acl_stream is the raw stream in Ascend NPU. Just like cuda_stream to GPU. 
What about name this stream torch_stream? since acl_stream and cuda_stream are both torch_stream now.
created_at: 2025-03-06T02:39:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,242 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+class _HcclGroup(Communicator):
+    """
+    Represents an actor's HCCL communicator. This is the default HCCL communicator
+    to be used in Compiled Graph if a custom communicator is not provided.
+
+    This class is not thread-safe.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        comm_id: int,
+        rank: Optional[int],
+        actor_handles: List["ray.actor.ActorHandle"],
+        acl_stream: Optional[Any],
+        use_communication_streams: bool = False,
+    ):
+        self._world_size = world_size
+        self._rank: Optional[int] = rank
+        self.hccl: Optional[ModuleType] = None
+        self._actor_handles = actor_handles
+        self._use_communication_streams = use_communication_streams
+
+        if rank is not None:
+            assert "NPU" in ray.cluster_resources(), "HCCL actor has no NPUs assigned"
+            assert acl_stream is not None, "HCCL actor must specify aclrtStream"
+
+            expected_rank = self.get_rank(ray.get_runtime_context().current_actor)
+            assert (
+                rank == expected_rank
+            ), f"HCCL actor's rank {rank} does not match expected rank {expected_rank}"
+
+            from ray.experimental.channel import hccl
+
+            self.hccl = hccl
+
+            self._comm = self.hccl.HCCLCommunicator(world_size, comm_id, rank)
+        else:
+            # Driver does not have a rank.
+            self._comm = None
+
+        self._acl_stream: None
+        self._send_stream: None
+        self._recv_stream: None
+        if acl_stream is not None:
+            assert rank is not None, "HCCL actor has no rank assigned"
+
+            self._acl_stream = acl_stream
+
+            if use_communication_streams:
+                import torch
+                import torch_npu  # noqa: F401
+                from ray.air._internal import torch_utils
+
+                # TODO(swang): Allow default device to be overridden.
+                device = torch_utils.get_devices()[0]
+
+                self._send_stream = torch.npu.Stream(device=device)
+                self._recv_stream = torch.npu.Stream(device=device)
+            else:
+                self._send_stream = self._acl_stream
+                self._recv_stream = self._acl_stream
+
+        self._closed = False
+
+    def initialize(self, rank: int) -> None:
+        pass
+
+    def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        return self._actor_handles
+
+    def get_rank(self, actor: ray.actor.ActorHandle) -> int:
+        actor_ids = [a._ray_actor_id for a in self._actor_handles]
+        try:
+            rank = actor_ids.index(actor._ray_actor_id)
+        except ValueError:
+            raise ValueError("Actor is not in the HCCL group.")
+        return rank
+
+    def get_self_rank(self) -> Optional[int]:
+        return self._rank
+
+    def get_world_size(self) -> int:
+        return self._world_size
+
+    def send(self, buf: "torch.Tensor", peer_rank: int) -> None:
path: python/ray/experimental/channel/hccl_group.py
user: login: hipudding
body: Yes, we should indeed reduce duplicate code. I will think about how to design it.
created_at: 2025-03-06T02:40:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -624,12 +623,12 @@ def wrap_and_set_intermediate_future(
         assert self._intermediate_future is None
 
         if wrap_in_gpu_future:
-            future = GPUFuture(val)
+            future = GPUFuture(val, stream)
         else:
             future = ResolvedFuture(val)
         self._intermediate_future = future
 
-    def reset_and_wait_intermediate_future(self) -> Any:
+    def reset_and_wait_intermediate_future(self, stream: Any = None) -> Any:
path: python/ray/dag/compiled_dag_node.py
user: login: hipudding
body: Done. reverted this part.
created_at: 2025-03-06T02:43:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -766,14 +767,12 @@ def exec_operation(
         """
         if op_type == _DAGNodeOperationType.READ:
             with _device_context_manager():
-                with self._recv_stream:
-                    return self._read(overlap_gpu_communication)
+                return self._read(overlap_gpu_communication)
         elif op_type == _DAGNodeOperationType.COMPUTE:
             return self._compute(overlap_gpu_communication, class_handle)
         elif op_type == _DAGNodeOperationType.WRITE:
             with _device_context_manager():
-                with self._send_stream:
-                    return self._write()
path: python/ray/dag/compiled_dag_node.py
user: login: hipudding
body: Done. use torch StreamContext instead to avoid too many changes.
created_at: 2025-03-06T07:29:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -600,15 +596,18 @@ def prepare(self, overlap_gpu_communication: bool = False):
                         nccl_group_id
                     )
                     assert nccl_group is not None
-                    if not isinstance(self._recv_stream, nullcontext):
+                    if self._recv_stream is not None:
                         assert self._recv_stream == nccl_group.recv_stream, (
                             "Currently all torch tensor input channels of a "
                             "Compiled Graph task should use the same recv cuda stream."
                         )
                     self._recv_stream = nccl_group.recv_stream
 
     def wrap_and_set_intermediate_future(
-        self, val: Any, wrap_in_gpu_future: bool
+        self,
+        val: Any,
+        wrap_in_gpu_future: bool,
+        stream: Any = None,
path: python/ray/dag/compiled_dag_node.py
user: login: hipudding
body: Done. reverted this part.
created_at: 2025-03-06T08:03:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,242 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+class _HcclGroup(Communicator):
+    """
+    Represents an actor's HCCL communicator. This is the default HCCL communicator
+    to be used in Compiled Graph if a custom communicator is not provided.
+
+    This class is not thread-safe.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        comm_id: int,
+        rank: Optional[int],
+        actor_handles: List["ray.actor.ActorHandle"],
+        acl_stream: Optional[Any],
+        use_communication_streams: bool = False,
+    ):
+        self._world_size = world_size
+        self._rank: Optional[int] = rank
+        self.hccl: Optional[ModuleType] = None
+        self._actor_handles = actor_handles
+        self._use_communication_streams = use_communication_streams
+
+        if rank is not None:
+            assert "NPU" in ray.cluster_resources(), "HCCL actor has no NPUs assigned"
+            assert acl_stream is not None, "HCCL actor must specify aclrtStream"
+
+            expected_rank = self.get_rank(ray.get_runtime_context().current_actor)
+            assert (
+                rank == expected_rank
+            ), f"HCCL actor's rank {rank} does not match expected rank {expected_rank}"
+
+            from ray.experimental.channel import hccl
+
+            self.hccl = hccl
+
+            self._comm = self.hccl.HCCLCommunicator(world_size, comm_id, rank)
+        else:
+            # Driver does not have a rank.
+            self._comm = None
+
+        self._acl_stream: None
+        self._send_stream: None
+        self._recv_stream: None
+        if acl_stream is not None:
+            assert rank is not None, "HCCL actor has no rank assigned"
+
+            self._acl_stream = acl_stream
+
+            if use_communication_streams:
+                import torch
+                import torch_npu  # noqa: F401
+                from ray.air._internal import torch_utils
+
+                # TODO(swang): Allow default device to be overridden.
+                device = torch_utils.get_devices()[0]
+
+                self._send_stream = torch.npu.Stream(device=device)
+                self._recv_stream = torch.npu.Stream(device=device)
+            else:
+                self._send_stream = self._acl_stream
+                self._recv_stream = self._acl_stream
+
+        self._closed = False
+
+    def initialize(self, rank: int) -> None:
+        pass
+
+    def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        return self._actor_handles
+
+    def get_rank(self, actor: ray.actor.ActorHandle) -> int:
+        actor_ids = [a._ray_actor_id for a in self._actor_handles]
+        try:
+            rank = actor_ids.index(actor._ray_actor_id)
+        except ValueError:
+            raise ValueError("Actor is not in the HCCL group.")
+        return rank
+
+    def get_self_rank(self) -> Optional[int]:
+        return self._rank
+
+    def get_world_size(self) -> int:
+        return self._world_size
+
+    def send(self, buf: "torch.Tensor", peer_rank: int) -> None:
path: python/ray/experimental/channel/hccl_group.py
user: login: hipudding
body: Indeed, these two files share logical similarities and contain a lot of duplicate code. However, there are several differences in the details. For example, torch.cuda.xxx needs to be replaced with torch.npu.xxx, hccl_group requires an additional StreamContext, and the set device operation must be performed before get_unique_id and creating the HCCL group. Therefore, I haven’t found a good way to merge these two files. Do you have any suggestions?
created_at: 2025-03-07T02:02:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,242 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+class _HcclGroup(Communicator):
+    """
+    Represents an actor's HCCL communicator. This is the default HCCL communicator
+    to be used in Compiled Graph if a custom communicator is not provided.
+
+    This class is not thread-safe.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        comm_id: int,
+        rank: Optional[int],
+        actor_handles: List["ray.actor.ActorHandle"],
+        acl_stream: Optional[Any],
path: python/ray/experimental/channel/hccl_group.py
user: login: hipudding
body: Done. I changed it to npu_stream for better understanding.
created_at: 2025-03-07T02:03:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,81 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class _DriverGroupHolder(Communicator):
+    """
+    Communicator place holder for Driver, Since driver may has no
+    Accelerators, TorchDeviceManager cannot be used.
+    """
path: python/ray/experimental/channel/driver_communicator_holder.py
user: login: hipudding
body: Comment updated.
created_at: 2025-03-07T02:04:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,401 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+    import torch_npu  # noqa: F401
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+def set_default_device():
+    """
+    Set the default NPU device for PyTorch.
+
+    This function initializes the NPU device by selecting the first available device
+    from `torch_utils.get_devices()`. It then sets the selected device as the default
+    NPU device using `torch.npu.set_device(device)`.
+
+    Returns:
+        torch.device: The selected NPU device.
+    """
+    import torch
+    import torch_npu  # noqa: F401
+    from ray.air._internal import torch_utils
+
+    device = torch_utils.get_devices()[0]
+    torch.npu.set_device(device)
+    return device
+
+
+class StreamContext:
+    """
+    Context-manager that selects a given stream.
+
+    Stream context, in which torch.npu.current_stream will be set to the value in the context. This class comes from pytorch because torch for NPU does not implement this context
path: python/ray/experimental/channel/hccl_group.py
user: login: ruisearch42
body: line break
created_at: 2025-03-07T23:27:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,401 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+    import torch_npu  # noqa: F401
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+def set_default_device():
+    """
+    Set the default NPU device for PyTorch.
+
+    This function initializes the NPU device by selecting the first available device
+    from `torch_utils.get_devices()`. It then sets the selected device as the default
+    NPU device using `torch.npu.set_device(device)`.
+
+    Returns:
+        torch.device: The selected NPU device.
+    """
+    import torch
+    import torch_npu  # noqa: F401
+    from ray.air._internal import torch_utils
+
+    device = torch_utils.get_devices()[0]
+    torch.npu.set_device(device)
+    return device
+
+
+class StreamContext:
+    """
+    Context-manager that selects a given stream.
+
+    Stream context, in which torch.npu.current_stream will be set to the value in the context. This class comes from pytorch because torch for NPU does not implement this context
+
+    Args:
+        Stream (Stream): selected stream. This manager is a no-op if it's
path: python/ray/experimental/channel/hccl_group.py
user: login: ruisearch42
body: ```suggestion
        stream (Stream): selected stream. This context manager is a no-op if it's
```
created_at: 2025-03-07T23:27:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,401 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+    import torch_npu  # noqa: F401
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+def set_default_device():
+    """
+    Set the default NPU device for PyTorch.
+
+    This function initializes the NPU device by selecting the first available device
+    from `torch_utils.get_devices()`. It then sets the selected device as the default
+    NPU device using `torch.npu.set_device(device)`.
+
+    Returns:
+        torch.device: The selected NPU device.
+    """
+    import torch
+    import torch_npu  # noqa: F401
+    from ray.air._internal import torch_utils
+
+    device = torch_utils.get_devices()[0]
+    torch.npu.set_device(device)
+    return device
+
+
+class StreamContext:
+    """
+    Context-manager that selects a given stream.
+
+    Stream context, in which torch.npu.current_stream will be set to the value in the context. This class comes from pytorch because torch for NPU does not implement this context
+
+    Args:
+        Stream (Stream): selected stream. This manager is a no-op if it's
+            ``None``.
+    note:: Streams are per-device.
+    """
+
+    cur_stream: Optional["torch.npu.Stream"]
+
+    def __init__(self, stream: Optional["torch.npu.Stream"]):
+        import torch
+        import torch_npu  # noqa: F401
+
+        self.stream = stream
+        self.idx = torch.npu.current_device()
+        if self.idx is None:
+            self.idx = -1
+
+        self.src_prev_stream = None
+        self.dst_prev_stream = None
path: python/ray/experimental/channel/hccl_group.py
user: login: ruisearch42
body: can you please keep variable type hints?
created_at: 2025-03-07T23:39:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,401 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+    import torch_npu  # noqa: F401
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+def set_default_device():
+    """
+    Set the default NPU device for PyTorch.
+
+    This function initializes the NPU device by selecting the first available device
+    from `torch_utils.get_devices()`. It then sets the selected device as the default
+    NPU device using `torch.npu.set_device(device)`.
+
+    Returns:
+        torch.device: The selected NPU device.
+    """
+    import torch
+    import torch_npu  # noqa: F401
+    from ray.air._internal import torch_utils
+
+    device = torch_utils.get_devices()[0]
+    torch.npu.set_device(device)
+    return device
+
+
+class StreamContext:
path: python/ray/experimental/channel/hccl_group.py
user: login: ruisearch42
body: Ideally this should be done at torch? How difficult is that?
created_at: 2025-03-07T23:40:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,364 @@
+import ctypes
+
+import torch
+
+hccl_lib = ctypes.CDLL("libhccl.so")
+
+HCCL_ROOT_INFO_BYTES = 4108
+
+
+class HcclRootInfo(ctypes.Structure):
+    _fields_ = [("internal", ctypes.c_byte * HCCL_ROOT_INFO_BYTES)]
+
+
+class HcclComm(ctypes.Structure):
+    pass
+
+
+class aclrtStream(ctypes.Structure):
+    pass
+
+
+# typedef enum {
+#     HCCL_SUCCESS = 0,               /**< success */
+#     HCCL_E_PARA = 1,                /**< parameter error */
+#     HCCL_E_PTR = 2,                 /**< empty pointer */
+#     HCCL_E_MEMORY = 3,              /**< memory error */
+#     HCCL_E_INTERNAL = 4,            /**< internal error */
+#     HCCL_E_NOT_SUPPORT = 5,         /**< not support feature */
+#     HCCL_E_NOT_FOUND = 6,           /**< not found specific resource */
+#     HCCL_E_UNAVAIL = 7,             /**< resource unavailable */
+#     HCCL_E_SYSCALL = 8,             /**< call system interface error */
+#     HCCL_E_TIMEOUT = 9,             /**< timeout */
+#     HCCL_E_OPEN_FILE_FAILURE = 10,  /**< open file fail */
+#     HCCL_E_TCP_CONNECT = 11,        /**< tcp connect fail */
+#     HCCL_E_ROCE_CONNECT = 12,       /**< roce connect fail */
+#     HCCL_E_TCP_TRANSFER = 13,       /**< tcp transfer fail */
+#     HCCL_E_ROCE_TRANSFER = 14,      /**< roce transfer fail */
+#     HCCL_E_RUNTIME = 15,            /**< call runtime api fail */
+#     HCCL_E_DRV = 16,                /**< call driver api fail */
+#     HCCL_E_PROFILING = 17,          /**< call profiling api fail */
+#     HCCL_E_CCE = 18,                /**< call cce api fail */
+#     HCCL_E_NETWORK = 19,            /**< call network api fail */
+#     HCCL_E_AGAIN = 20,              /**< try again */
+#     HCCL_E_REMOTE = 21,             /**< error cqe */
+#     HCCL_E_SUSPENDING = 22,         /**< error communicator suspending */
+#     HCCL_E_RESERVED                 /**< reserved */
+# } HcclResult;
+
+HCCL_ERR_STR = {
+    0: "HCCL_SUCCESS",
+    1: "HCCL_E_PARA",
+    2: "HCCL_E_PTR",
+    3: "HCCL_E_MEMORY",
+    4: "HCCL_E_INTERNAL",
+    5: "HCCL_E_NOT_SUPPORT",
+    6: "HCCL_E_NOT_FOUND",
+    7: "HCCL_E_UNAVAIL",
+    8: "HCCL_E_SYSCALL",
+    9: "HCCL_E_TIMEOUT",
+    10: "HCCL_E_OPEN_FILE_FAILURE",
+    11: "HCCL_E_TCP_CONNECT",
+    12: "HCCL_E_ROCE_CONNECT",
+    13: "HCCL_E_TCP_TRANSFER",
+    14: "HCCL_E_ROCE_TRANSFER",
+    15: "HCCL_E_RUNTIME",
+    16: "HCCL_E_DRV",
+    17: "HCCL_E_PROFILING",
+    18: "HCCL_E_CCE",
+    19: "HCCL_E_NETWORK",
+    20: "HCCL_E_AGAIN",
+    21: "HCCL_E_REMOTE",
+    22: "HCCL_E_SUSPENDING",
+    23: "HCCL_E_RESERVED",
+}
+
+# typedef enum {
+#     HCCL_DATA_TYPE_INT8 = 0,    /**< int8 */
+#     HCCL_DATA_TYPE_INT16 = 1,   /**< int16 */
+#     HCCL_DATA_TYPE_INT32 = 2,   /**< int32 */
+#     HCCL_DATA_TYPE_FP16 = 3,    /**< fp16 */
+#     HCCL_DATA_TYPE_FP32 = 4,    /**< fp32 */
+#     HCCL_DATA_TYPE_INT64 = 5,    /**< int64 */
+#     HCCL_DATA_TYPE_UINT64 = 6,    /**< uint64 */
+#     HCCL_DATA_TYPE_UINT8 = 7,    /**< uint8 */
+#     HCCL_DATA_TYPE_UINT16 = 8,   /**< uint16 */
+#     HCCL_DATA_TYPE_UINT32 = 9,   /**< uint32 */
+#     HCCL_DATA_TYPE_FP64 = 10, /**< fp64 */
+#     HCCL_DATA_TYPE_BFP16 = 11,    /**< bfp16 */
+#     HCCL_DATA_TYPE_INT128 = 12,   /**< int128 */
+#     HCCL_DATA_TYPE_RESERVED     /**< reserved */
+# } HcclDataType;
+
+TORCH_HCCL_DTYPE_MAP = {
+    torch.bool: 0,  # HCCL_DATA_TYPE_INT8
+    # INT types
+    torch.int: 2,  # HCCL_DATA_TYPE_INT32
+    torch.uint8: 7,  # HCCL_DATA_TYPE_UINT8
+    torch.int8: 0,  # HCCL_DATA_TYPE_INT8
+    torch.int32: 2,  # HCCL_DATA_TYPE_INT32
+    torch.int64: 5,  # HCCL_DATA_TYPE_INT64
+    torch.long: 5,  # HCCL_DATA_TYPE_INT64
+    # FLOAT types
+    torch.half: 3,  # HCCL_DATA_TYPE_FP16
+    torch.float: 4,  # HCCL_DATA_TYPE_FP32
+    torch.float16: 3,  # HCCL_DATA_TYPE_FP16
+    torch.float32: 4,  # HCCL_DATA_TYPE_FP32
+    torch.float64: 10,  # HCCL_DATA_TYPE_FP64
+    torch.double: 10,  # HCCL_DATA_TYPE_FP64
+    torch.bfloat16: 11,  # HCCL_DATA_TYPE_BFP16
+}
+
+# typedef enum {
+#     HCCL_REDUCE_SUM = 0,    /**< sum */
+#     HCCL_REDUCE_PROD = 1,   /**< prod */
+#     HCCL_REDUCE_MAX = 2,    /**< max */
+#     HCCL_REDUCE_MIN = 3,    /**< min */
+#     HCCL_REDUCE_RESERVED    /**< reserved */
+# } HcclReduceOp;
+
+# HcclResult HcclGetRootInfo(HcclRootInfo *rootInfo);
+hccl_lib.HcclGetRootInfo.argtypes = [ctypes.POINTER(HcclRootInfo)]
+hccl_lib.HcclGetRootInfo.restype = ctypes.c_int
+
+# HcclResult HcclCommInitRootInfo(uint32_t nRanks, const HcclRootInfo *rootInfo,
+#                                   uint32_t rank, HcclComm *comm);
+hccl_lib.HcclCommInitRootInfo.argtypes = [
+    ctypes.c_uint32,
+    ctypes.POINTER(HcclRootInfo),
+    ctypes.c_uint32,
+    ctypes.POINTER(ctypes.POINTER(HcclComm)),
+]
+hccl_lib.HcclCommInitRootInfo.restype = ctypes.c_int
+
+# HcclResult HcclCommDestroy(HcclComm comm);
+hccl_lib.HcclCommDestroy.argtypes = [ctypes.POINTER(HcclComm)]
+hccl_lib.HcclCommDestroy.restype = ctypes.c_int
+
+# HcclResult HcclSend(void* sendBuf, uint64_t count, HcclDataType dataType,
+#                       uint32_t destRank, HcclComm comm, aclrtStream stream);
+hccl_lib.HcclSend.argtypes = [
+    ctypes.c_void_p,
+    ctypes.c_uint64,
+    ctypes.c_int,
+    ctypes.c_uint32,
+    ctypes.POINTER(HcclComm),
+    ctypes.POINTER(aclrtStream),
+]
+hccl_lib.HcclSend.restype = ctypes.c_int
+
+# HcclResult HcclRecv(void* recvBuf, uint64_t count, HcclDataType dataType,
+#                       uint32_t srcRank, HcclComm comm, aclrtStream stream);
+hccl_lib.HcclRecv.argtypes = [
+    ctypes.c_void_p,
+    ctypes.c_uint64,
+    ctypes.c_int,
+    ctypes.c_uint32,
+    ctypes.POINTER(HcclComm),
+    ctypes.POINTER(aclrtStream),
+]
+hccl_lib.HcclRecv.restype = ctypes.c_int
+
+# HcclResult HcclAllReduce(void *sendBuf, void *recvBuf, uint64_t count,
+#                           HcclDataType dataType,  HcclReduceOp op,
+#                           HcclComm comm, aclrtStream stream);
+hccl_lib.HcclAllReduce.argtypes = [
+    ctypes.c_void_p,
+    ctypes.c_void_p,
+    ctypes.c_uint64,
+    ctypes.c_int,
+    ctypes.c_int,
+    ctypes.POINTER(HcclComm),
+    ctypes.POINTER(aclrtStream),
+]
+hccl_lib.HcclAllReduce.restype = ctypes.c_int
+
+# HcclResult HcclReduce(void *sendBuf, void *recvBuf, uint64_t count,
+#                       HcclDataType dataType, HcclReduceOp op,
+#                       uint32_t root, HcclComm comm, aclrtStream stream);
+hccl_lib.HcclReduce.argtypes = [
+    ctypes.c_void_p,
+    ctypes.c_void_p,
+    ctypes.c_uint64,
+    ctypes.c_int,
+    ctypes.c_int,
+    ctypes.c_uint32,
+    ctypes.POINTER(HcclComm),
+    ctypes.POINTER(aclrtStream),
+]
+hccl_lib.HcclReduce.restype = ctypes.c_int
+
+# HcclResult HcclBroadcast(void *buf, uint64_t count,
+#                           HcclDataType dataType, uint32_t root,
+#                           HcclComm comm, aclrtStream stream);
+hccl_lib.HcclBroadcast.argtypes = [
+    ctypes.c_void_p,
+    ctypes.c_uint64,
+    ctypes.c_int,
+    ctypes.c_uint32,
+    ctypes.POINTER(HcclComm),
+    ctypes.POINTER(aclrtStream),
+]
+hccl_lib.HcclBroadcast.restype = ctypes.c_int
+
+# HcclResult HcclReduceScatter(void *sendBuf, void *recvBuf, uint64_t recvCount,
+#                               HcclDataType dataType, HcclReduceOp op,
+#                               HcclComm comm, aclrtStream stream);
+hccl_lib.HcclReduceScatter.argtypes = [
+    ctypes.c_void_p,
+    ctypes.c_void_p,
+    ctypes.c_uint64,
+    ctypes.c_int,
+    ctypes.c_int,
+    ctypes.POINTER(HcclComm),
+    ctypes.POINTER(aclrtStream),
+]
+hccl_lib.HcclReduceScatter.restype = ctypes.c_int
+
+# HcclResult HcclAllReduce(void *sendBuf, void *recvBuf, uint64_t count,
+#                           HcclDataType dataType, HcclReduceOp op,
+#                           HcclComm comm, aclrtStream stream);
+hccl_lib.HcclAllReduce.argtypes = [
+    ctypes.c_void_p,
+    ctypes.c_void_p,
+    ctypes.c_uint64,
+    ctypes.c_int,
+    ctypes.c_int,
+    ctypes.POINTER(HcclComm),
+    ctypes.POINTER(aclrtStream),
+]
+hccl_lib.HcclAllReduce.restype = ctypes.c_int
+
+# const char *HcclGetErrorString(HcclResult code);
+hccl_lib.HcclGetErrorString.argtypes = [ctypes.c_int]
+hccl_lib.HcclGetErrorString.restype = ctypes.c_char_p
+
+# HcclResult enum (simplified)
+HCCL_SUCCESS = 0
+
+
+class HcclError(RuntimeError):
+    def __init__(self, status):
+        self.status = status
+        error_message = hccl_lib.HcclGetErrorString(status).decode()
+        super(HcclError, self).__init__(
+            "%s: %s" % (HCCL_ERR_STR[status], error_message)
+        )
+
+    def __reduce__(self):
+        return (type(self), (self.status,))
+
+
+def check_hccl_status(status):
+    if status != HCCL_SUCCESS:
+        raise HcclError(status)
+
+
+def get_unique_id():
+    root_info = HcclRootInfo()
+    status = hccl_lib.HcclGetRootInfo(root_info)
+    check_hccl_status(status)
+    ret = tuple([root_info.internal[i] for i in range(HCCL_ROOT_INFO_BYTES)])
+    return ret
+
+
+def get_hccl_tensor_dtype(tensor):
+    return TORCH_HCCL_DTYPE_MAP[tensor.dtype]
+
+
+class HCCLCommunicator:
path: python/ray/experimental/channel/hccl.py
user: login: ruisearch42
body: How did you do testing? Not sure if we have NPUs for the CI. And we will need some testing strategy
created_at: 2025-03-07T23:54:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,401 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+    import torch_npu  # noqa: F401
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+def set_default_device():
+    """
+    Set the default NPU device for PyTorch.
+
+    This function initializes the NPU device by selecting the first available device
+    from `torch_utils.get_devices()`. It then sets the selected device as the default
+    NPU device using `torch.npu.set_device(device)`.
+
+    Returns:
+        torch.device: The selected NPU device.
+    """
+    import torch
+    import torch_npu  # noqa: F401
+    from ray.air._internal import torch_utils
+
+    device = torch_utils.get_devices()[0]
+    torch.npu.set_device(device)
+    return device
+
+
+class StreamContext:
+    """
+    Context-manager that selects a given stream.
+
+    Stream context, in which torch.npu.current_stream will be set to the value in the context. This class comes from pytorch because torch for NPU does not implement this context
path: python/ray/experimental/channel/hccl_group.py
user: login: hipudding
body: Done.
created_at: 2025-03-10T01:06:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,401 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+    import torch_npu  # noqa: F401
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+def set_default_device():
+    """
+    Set the default NPU device for PyTorch.
+
+    This function initializes the NPU device by selecting the first available device
+    from `torch_utils.get_devices()`. It then sets the selected device as the default
+    NPU device using `torch.npu.set_device(device)`.
+
+    Returns:
+        torch.device: The selected NPU device.
+    """
+    import torch
+    import torch_npu  # noqa: F401
+    from ray.air._internal import torch_utils
+
+    device = torch_utils.get_devices()[0]
+    torch.npu.set_device(device)
+    return device
+
+
+class StreamContext:
+    """
+    Context-manager that selects a given stream.
+
+    Stream context, in which torch.npu.current_stream will be set to the value in the context. This class comes from pytorch because torch for NPU does not implement this context
+
+    Args:
+        Stream (Stream): selected stream. This manager is a no-op if it's
path: python/ray/experimental/channel/hccl_group.py
user: login: hipudding
body: Done.
created_at: 2025-03-10T01:06:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,401 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+    import torch_npu  # noqa: F401
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+def set_default_device():
+    """
+    Set the default NPU device for PyTorch.
+
+    This function initializes the NPU device by selecting the first available device
+    from `torch_utils.get_devices()`. It then sets the selected device as the default
+    NPU device using `torch.npu.set_device(device)`.
+
+    Returns:
+        torch.device: The selected NPU device.
+    """
+    import torch
+    import torch_npu  # noqa: F401
+    from ray.air._internal import torch_utils
+
+    device = torch_utils.get_devices()[0]
+    torch.npu.set_device(device)
+    return device
+
+
+class StreamContext:
+    """
+    Context-manager that selects a given stream.
+
+    Stream context, in which torch.npu.current_stream will be set to the value in the context. This class comes from pytorch because torch for NPU does not implement this context
+
+    Args:
+        Stream (Stream): selected stream. This manager is a no-op if it's
+            ``None``.
+    note:: Streams are per-device.
+    """
+
+    cur_stream: Optional["torch.npu.Stream"]
+
+    def __init__(self, stream: Optional["torch.npu.Stream"]):
+        import torch
+        import torch_npu  # noqa: F401
+
+        self.stream = stream
+        self.idx = torch.npu.current_device()
+        if self.idx is None:
+            self.idx = -1
+
+        self.src_prev_stream = None
+        self.dst_prev_stream = None
path: python/ray/experimental/channel/hccl_group.py
user: login: hipudding
body: Done.
created_at: 2025-03-10T01:06:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,401 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+    import torch_npu  # noqa: F401
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+def set_default_device():
+    """
+    Set the default NPU device for PyTorch.
+
+    This function initializes the NPU device by selecting the first available device
+    from `torch_utils.get_devices()`. It then sets the selected device as the default
+    NPU device using `torch.npu.set_device(device)`.
+
+    Returns:
+        torch.device: The selected NPU device.
+    """
+    import torch
+    import torch_npu  # noqa: F401
+    from ray.air._internal import torch_utils
+
+    device = torch_utils.get_devices()[0]
+    torch.npu.set_device(device)
+    return device
+
+
+class StreamContext:
path: python/ray/experimental/channel/hccl_group.py
user: login: hipudding
body: Yes, the StreamContext should indeed be implemented in Torch. However, the torch_npu team has not yet implemented this structure. I have raised the request with them and hope they will support it soon. For now, to maintain compatibility with older versions of Torch, this StreamContext is still necessary. Once torch_npu supports it, this StreamContext will be removed.
created_at: 2025-03-10T01:08:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,364 @@
+import ctypes
+
+import torch
+
+hccl_lib = ctypes.CDLL("libhccl.so")
+
+HCCL_ROOT_INFO_BYTES = 4108
+
+
+class HcclRootInfo(ctypes.Structure):
+    _fields_ = [("internal", ctypes.c_byte * HCCL_ROOT_INFO_BYTES)]
+
+
+class HcclComm(ctypes.Structure):
+    pass
+
+
+class aclrtStream(ctypes.Structure):
+    pass
+
+
+# typedef enum {
+#     HCCL_SUCCESS = 0,               /**< success */
+#     HCCL_E_PARA = 1,                /**< parameter error */
+#     HCCL_E_PTR = 2,                 /**< empty pointer */
+#     HCCL_E_MEMORY = 3,              /**< memory error */
+#     HCCL_E_INTERNAL = 4,            /**< internal error */
+#     HCCL_E_NOT_SUPPORT = 5,         /**< not support feature */
+#     HCCL_E_NOT_FOUND = 6,           /**< not found specific resource */
+#     HCCL_E_UNAVAIL = 7,             /**< resource unavailable */
+#     HCCL_E_SYSCALL = 8,             /**< call system interface error */
+#     HCCL_E_TIMEOUT = 9,             /**< timeout */
+#     HCCL_E_OPEN_FILE_FAILURE = 10,  /**< open file fail */
+#     HCCL_E_TCP_CONNECT = 11,        /**< tcp connect fail */
+#     HCCL_E_ROCE_CONNECT = 12,       /**< roce connect fail */
+#     HCCL_E_TCP_TRANSFER = 13,       /**< tcp transfer fail */
+#     HCCL_E_ROCE_TRANSFER = 14,      /**< roce transfer fail */
+#     HCCL_E_RUNTIME = 15,            /**< call runtime api fail */
+#     HCCL_E_DRV = 16,                /**< call driver api fail */
+#     HCCL_E_PROFILING = 17,          /**< call profiling api fail */
+#     HCCL_E_CCE = 18,                /**< call cce api fail */
+#     HCCL_E_NETWORK = 19,            /**< call network api fail */
+#     HCCL_E_AGAIN = 20,              /**< try again */
+#     HCCL_E_REMOTE = 21,             /**< error cqe */
+#     HCCL_E_SUSPENDING = 22,         /**< error communicator suspending */
+#     HCCL_E_RESERVED                 /**< reserved */
+# } HcclResult;
+
+HCCL_ERR_STR = {
+    0: "HCCL_SUCCESS",
+    1: "HCCL_E_PARA",
+    2: "HCCL_E_PTR",
+    3: "HCCL_E_MEMORY",
+    4: "HCCL_E_INTERNAL",
+    5: "HCCL_E_NOT_SUPPORT",
+    6: "HCCL_E_NOT_FOUND",
+    7: "HCCL_E_UNAVAIL",
+    8: "HCCL_E_SYSCALL",
+    9: "HCCL_E_TIMEOUT",
+    10: "HCCL_E_OPEN_FILE_FAILURE",
+    11: "HCCL_E_TCP_CONNECT",
+    12: "HCCL_E_ROCE_CONNECT",
+    13: "HCCL_E_TCP_TRANSFER",
+    14: "HCCL_E_ROCE_TRANSFER",
+    15: "HCCL_E_RUNTIME",
+    16: "HCCL_E_DRV",
+    17: "HCCL_E_PROFILING",
+    18: "HCCL_E_CCE",
+    19: "HCCL_E_NETWORK",
+    20: "HCCL_E_AGAIN",
+    21: "HCCL_E_REMOTE",
+    22: "HCCL_E_SUSPENDING",
+    23: "HCCL_E_RESERVED",
+}
+
+# typedef enum {
+#     HCCL_DATA_TYPE_INT8 = 0,    /**< int8 */
+#     HCCL_DATA_TYPE_INT16 = 1,   /**< int16 */
+#     HCCL_DATA_TYPE_INT32 = 2,   /**< int32 */
+#     HCCL_DATA_TYPE_FP16 = 3,    /**< fp16 */
+#     HCCL_DATA_TYPE_FP32 = 4,    /**< fp32 */
+#     HCCL_DATA_TYPE_INT64 = 5,    /**< int64 */
+#     HCCL_DATA_TYPE_UINT64 = 6,    /**< uint64 */
+#     HCCL_DATA_TYPE_UINT8 = 7,    /**< uint8 */
+#     HCCL_DATA_TYPE_UINT16 = 8,   /**< uint16 */
+#     HCCL_DATA_TYPE_UINT32 = 9,   /**< uint32 */
+#     HCCL_DATA_TYPE_FP64 = 10, /**< fp64 */
+#     HCCL_DATA_TYPE_BFP16 = 11,    /**< bfp16 */
+#     HCCL_DATA_TYPE_INT128 = 12,   /**< int128 */
+#     HCCL_DATA_TYPE_RESERVED     /**< reserved */
+# } HcclDataType;
+
+TORCH_HCCL_DTYPE_MAP = {
+    torch.bool: 0,  # HCCL_DATA_TYPE_INT8
+    # INT types
+    torch.int: 2,  # HCCL_DATA_TYPE_INT32
+    torch.uint8: 7,  # HCCL_DATA_TYPE_UINT8
+    torch.int8: 0,  # HCCL_DATA_TYPE_INT8
+    torch.int32: 2,  # HCCL_DATA_TYPE_INT32
+    torch.int64: 5,  # HCCL_DATA_TYPE_INT64
+    torch.long: 5,  # HCCL_DATA_TYPE_INT64
+    # FLOAT types
+    torch.half: 3,  # HCCL_DATA_TYPE_FP16
+    torch.float: 4,  # HCCL_DATA_TYPE_FP32
+    torch.float16: 3,  # HCCL_DATA_TYPE_FP16
+    torch.float32: 4,  # HCCL_DATA_TYPE_FP32
+    torch.float64: 10,  # HCCL_DATA_TYPE_FP64
+    torch.double: 10,  # HCCL_DATA_TYPE_FP64
+    torch.bfloat16: 11,  # HCCL_DATA_TYPE_BFP16
+}
+
+# typedef enum {
+#     HCCL_REDUCE_SUM = 0,    /**< sum */
+#     HCCL_REDUCE_PROD = 1,   /**< prod */
+#     HCCL_REDUCE_MAX = 2,    /**< max */
+#     HCCL_REDUCE_MIN = 3,    /**< min */
+#     HCCL_REDUCE_RESERVED    /**< reserved */
+# } HcclReduceOp;
+
+# HcclResult HcclGetRootInfo(HcclRootInfo *rootInfo);
+hccl_lib.HcclGetRootInfo.argtypes = [ctypes.POINTER(HcclRootInfo)]
+hccl_lib.HcclGetRootInfo.restype = ctypes.c_int
+
+# HcclResult HcclCommInitRootInfo(uint32_t nRanks, const HcclRootInfo *rootInfo,
+#                                   uint32_t rank, HcclComm *comm);
+hccl_lib.HcclCommInitRootInfo.argtypes = [
+    ctypes.c_uint32,
+    ctypes.POINTER(HcclRootInfo),
+    ctypes.c_uint32,
+    ctypes.POINTER(ctypes.POINTER(HcclComm)),
+]
+hccl_lib.HcclCommInitRootInfo.restype = ctypes.c_int
+
+# HcclResult HcclCommDestroy(HcclComm comm);
+hccl_lib.HcclCommDestroy.argtypes = [ctypes.POINTER(HcclComm)]
+hccl_lib.HcclCommDestroy.restype = ctypes.c_int
+
+# HcclResult HcclSend(void* sendBuf, uint64_t count, HcclDataType dataType,
+#                       uint32_t destRank, HcclComm comm, aclrtStream stream);
+hccl_lib.HcclSend.argtypes = [
+    ctypes.c_void_p,
+    ctypes.c_uint64,
+    ctypes.c_int,
+    ctypes.c_uint32,
+    ctypes.POINTER(HcclComm),
+    ctypes.POINTER(aclrtStream),
+]
+hccl_lib.HcclSend.restype = ctypes.c_int
+
+# HcclResult HcclRecv(void* recvBuf, uint64_t count, HcclDataType dataType,
+#                       uint32_t srcRank, HcclComm comm, aclrtStream stream);
+hccl_lib.HcclRecv.argtypes = [
+    ctypes.c_void_p,
+    ctypes.c_uint64,
+    ctypes.c_int,
+    ctypes.c_uint32,
+    ctypes.POINTER(HcclComm),
+    ctypes.POINTER(aclrtStream),
+]
+hccl_lib.HcclRecv.restype = ctypes.c_int
+
+# HcclResult HcclAllReduce(void *sendBuf, void *recvBuf, uint64_t count,
+#                           HcclDataType dataType,  HcclReduceOp op,
+#                           HcclComm comm, aclrtStream stream);
+hccl_lib.HcclAllReduce.argtypes = [
+    ctypes.c_void_p,
+    ctypes.c_void_p,
+    ctypes.c_uint64,
+    ctypes.c_int,
+    ctypes.c_int,
+    ctypes.POINTER(HcclComm),
+    ctypes.POINTER(aclrtStream),
+]
+hccl_lib.HcclAllReduce.restype = ctypes.c_int
+
+# HcclResult HcclReduce(void *sendBuf, void *recvBuf, uint64_t count,
+#                       HcclDataType dataType, HcclReduceOp op,
+#                       uint32_t root, HcclComm comm, aclrtStream stream);
+hccl_lib.HcclReduce.argtypes = [
+    ctypes.c_void_p,
+    ctypes.c_void_p,
+    ctypes.c_uint64,
+    ctypes.c_int,
+    ctypes.c_int,
+    ctypes.c_uint32,
+    ctypes.POINTER(HcclComm),
+    ctypes.POINTER(aclrtStream),
+]
+hccl_lib.HcclReduce.restype = ctypes.c_int
+
+# HcclResult HcclBroadcast(void *buf, uint64_t count,
+#                           HcclDataType dataType, uint32_t root,
+#                           HcclComm comm, aclrtStream stream);
+hccl_lib.HcclBroadcast.argtypes = [
+    ctypes.c_void_p,
+    ctypes.c_uint64,
+    ctypes.c_int,
+    ctypes.c_uint32,
+    ctypes.POINTER(HcclComm),
+    ctypes.POINTER(aclrtStream),
+]
+hccl_lib.HcclBroadcast.restype = ctypes.c_int
+
+# HcclResult HcclReduceScatter(void *sendBuf, void *recvBuf, uint64_t recvCount,
+#                               HcclDataType dataType, HcclReduceOp op,
+#                               HcclComm comm, aclrtStream stream);
+hccl_lib.HcclReduceScatter.argtypes = [
+    ctypes.c_void_p,
+    ctypes.c_void_p,
+    ctypes.c_uint64,
+    ctypes.c_int,
+    ctypes.c_int,
+    ctypes.POINTER(HcclComm),
+    ctypes.POINTER(aclrtStream),
+]
+hccl_lib.HcclReduceScatter.restype = ctypes.c_int
+
+# HcclResult HcclAllReduce(void *sendBuf, void *recvBuf, uint64_t count,
+#                           HcclDataType dataType, HcclReduceOp op,
+#                           HcclComm comm, aclrtStream stream);
+hccl_lib.HcclAllReduce.argtypes = [
+    ctypes.c_void_p,
+    ctypes.c_void_p,
+    ctypes.c_uint64,
+    ctypes.c_int,
+    ctypes.c_int,
+    ctypes.POINTER(HcclComm),
+    ctypes.POINTER(aclrtStream),
+]
+hccl_lib.HcclAllReduce.restype = ctypes.c_int
+
+# const char *HcclGetErrorString(HcclResult code);
+hccl_lib.HcclGetErrorString.argtypes = [ctypes.c_int]
+hccl_lib.HcclGetErrorString.restype = ctypes.c_char_p
+
+# HcclResult enum (simplified)
+HCCL_SUCCESS = 0
+
+
+class HcclError(RuntimeError):
+    def __init__(self, status):
+        self.status = status
+        error_message = hccl_lib.HcclGetErrorString(status).decode()
+        super(HcclError, self).__init__(
+            "%s: %s" % (HCCL_ERR_STR[status], error_message)
+        )
+
+    def __reduce__(self):
+        return (type(self), (self.status,))
+
+
+def check_hccl_status(status):
+    if status != HCCL_SUCCESS:
+        raise HcclError(status)
+
+
+def get_unique_id():
+    root_info = HcclRootInfo()
+    status = hccl_lib.HcclGetRootInfo(root_info)
+    check_hccl_status(status)
+    ret = tuple([root_info.internal[i] for i in range(HCCL_ROOT_INFO_BYTES)])
+    return ret
+
+
+def get_hccl_tensor_dtype(tensor):
+    return TORCH_HCCL_DTYPE_MAP[tensor.dtype]
+
+
+class HCCLCommunicator:
path: python/ray/experimental/channel/hccl.py
user: login: hipudding
body: I indirectly verified the correctness of HCCLCommunicator in the NPU environment through the test cases in test_torch_tensor_dag.py. However, since the Ray test environment does not have Ascend devices, these tests cannot currently be validated in the workflow. We are more than happy to contribute the Ascend NPU CI to Ray’s infrastructure to facilitate the validation of NPU-related functionalities.
created_at: 2025-03-10T01:10:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,401 @@
+import logging
+from types import ModuleType
+from typing import TYPE_CHECKING, List, Optional, Tuple, Any
+
+import ray
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+    import torch_npu  # noqa: F401
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+
+def set_default_device():
+    """
+    Set the default NPU device for PyTorch.
+
+    This function initializes the NPU device by selecting the first available device
+    from `torch_utils.get_devices()`. It then sets the selected device as the default
+    NPU device using `torch.npu.set_device(device)`.
+
+    Returns:
+        torch.device: The selected NPU device.
+    """
+    import torch
+    import torch_npu  # noqa: F401
+    from ray.air._internal import torch_utils
+
+    device = torch_utils.get_devices()[0]
+    torch.npu.set_device(device)
+    return device
+
+
+class StreamContext:
path: python/ray/experimental/channel/hccl_group.py
user: login: ruisearch42
body: Thanks. Can you link the corresponding GH issue and add a TODO here?
created_at: 2025-03-11T20:42:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -93,29 +91,23 @@ def __init__(
             # Driver does not have a rank.
             self._comm = None
 
-        self._cuda_stream: Optional["cp.cuda.ExternalStream"] = None
-        self._send_stream: Optional["cp.cuda.ExternalStream"] = None
-        self._recv_stream: Optional["cp.cuda.ExternalStream"] = None
+        self._cuda_stream: Optional["torch.cuda.Stream"] = None
+        self._send_stream: Optional["torch.cuda.Stream"] = None
+        self._recv_stream: Optional["torch.cuda.Stream"] = None
         if cuda_stream is not None:
             assert rank is not None, "NCCL actor has no rank assigned"
 
-            import cupy as cp
-
-            # TODO(swang): Allow default device to be overridden.
-            device = get_default_torch_device(allow_cpu=False)
-            self._cuda_stream = cp.cuda.ExternalStream(
-                cuda_stream, device_id=device.index
-            )
+            self._cuda_stream = cuda_stream
 
             if use_communication_streams:
                 import torch
+                from ray.air._internal import torch_utils
 
-                self._send_stream = cp.cuda.ExternalStream(
-                    torch.cuda.Stream().cuda_stream, device_id=device.index
-                )
-                self._recv_stream = cp.cuda.ExternalStream(
-                    torch.cuda.Stream().cuda_stream, device_id=device.index
-                )
+                # TODO(swang): Allow default device to be overridden.
+                device = torch_utils.get_devices()[0]
path: python/ray/experimental/channel/nccl_group.py
user: login: hipudding
body: ```suggestion
                device = get_default_torch_device(allow_cpu=False)
```
created_at: 2025-03-18T12:28:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -159,11 +159,13 @@ def deserialize_from_numpy(
             target_device_type = tensor_device_type
         elif target_device in [Device.GPU, Device.CUDA]:
             target_device_type = "cuda"
+        elif target_device == Device.NPU:
+            target_device_type = "npu"
         else:
             target_device_type = "cpu"
 
         # TODO(swang): Support local P2P transfers if available.
-        if target_device_type == "cuda":
+        if target_device_type in ["cuda", "npu"]:
path: python/ray/experimental/channel/serialization_context.py
user: login: hipudding
body: Use better device capability check.
created_at: 2025-03-18T12:30:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -448,13 +448,13 @@ def ensure_registered_as_writer(self):
         assert self._nccl_group is not None, "Actor is not part of a NCCL group"
         assert self._writer_registered
         ctx = ChannelContext.get_current()
-        assert ctx.torch_device.type == "cuda"
+        assert ctx.torch_device.type in ["cuda", "npu"]
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: hipudding
body: Use better device capability check.
created_at: 2025-03-18T12:30:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -448,13 +448,13 @@ def ensure_registered_as_writer(self):
         assert self._nccl_group is not None, "Actor is not part of a NCCL group"
         assert self._writer_registered
         ctx = ChannelContext.get_current()
-        assert ctx.torch_device.type == "cuda"
+        assert ctx.torch_device.type in ["cuda", "npu"]
 
     def ensure_registered_as_reader(self) -> bool:
         assert self._nccl_group is not None, "Actor is not part of a NCCL group"
         assert self._reader_registered
         ctx = ChannelContext.get_current()
-        assert ctx.torch_device.type == "cuda"
+        assert ctx.torch_device.type in ["cuda", "npu"]
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: hipudding
body: Use better device capability check.
created_at: 2025-03-18T12:30:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -591,8 +590,8 @@ def prepare(self, overlap_gpu_communication: bool = False):
         self.input_reader.start()
         self.output_writer.start()
 
-        self._send_stream: Union["cp.cuda.Stream", nullcontext] = nullcontext()
-        self._recv_stream: Union["cp.cuda.Stream", nullcontext] = nullcontext()
path: python/ray/dag/compiled_dag_node.py
user: login: ruisearch42
body: Can we keep the type hints?
created_at: 2025-03-19T05:24:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -191,7 +191,7 @@ def with_tensor_transport(
         else:
             if not isinstance(transport, Communicator):
                 raise ValueError(
-                    "transport must be 'auto', 'nccl' or a Communicator type"
+                    "transport must be 'auto', 'nccl', 'cpu' or a Communicator type"
path: python/ray/dag/dag_node.py
user: login: ruisearch42
body: This is not consistent with L176
created_at: 2025-03-19T05:26:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -184,3 +185,15 @@ def recv_stream(self):
 
     def send_stream(self):
         raise NotImplementedError
+
+
+def get_cpu_unique_id():
path: python/ray/experimental/channel/cpu_communicator.py
user: login: ruisearch42
body: `generate_uid_id()` is better?
created_at: 2025-03-19T05:27:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,91 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class _DriverGroupHolder(Communicator):
path: python/ray/experimental/channel/driver_communicator_holder.py
user: login: ruisearch42
body: probably we could have a better name, something like `CommunicatorInfoHolder` but better than it
created_at: 2025-03-19T05:30:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,364 @@
+import ctypes
+
+import torch
+
+hccl_lib = ctypes.CDLL("libhccl.so")
+
+HCCL_ROOT_INFO_BYTES = 4108
+
+
+class HcclRootInfo(ctypes.Structure):
+    _fields_ = [("internal", ctypes.c_byte * HCCL_ROOT_INFO_BYTES)]
+
+
+class HcclComm(ctypes.Structure):
+    pass
+
+
+class aclrtStream(ctypes.Structure):
path: python/ray/experimental/channel/hccl.py
user: login: ruisearch42
body: docs
created_at: 2025-03-19T05:30:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -591,8 +590,8 @@ def prepare(self, overlap_gpu_communication: bool = False):
         self.input_reader.start()
         self.output_writer.start()
 
-        self._send_stream: Union["cp.cuda.Stream", nullcontext] = nullcontext()
-        self._recv_stream: Union["cp.cuda.Stream", nullcontext] = nullcontext()
path: python/ray/dag/compiled_dag_node.py
user: login: hipudding
body: I'm afraid we cannot keep this type hints. For cuda, the stream type is torch.cuda.Stream, for NPU the stream type is torch.npu.Stream. Torch 2.6 support a general type torch.Stream, But I think we should support old torch versions.
created_at: 2025-03-19T07:00:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -184,3 +185,15 @@ def recv_stream(self):
 
     def send_stream(self):
         raise NotImplementedError
+
+
+def get_cpu_unique_id():
path: python/ray/experimental/channel/cpu_communicator.py
user: login: hipudding
body: Done.
created_at: 2025-03-19T07:01:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,91 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class _DriverGroupHolder(Communicator):
path: python/ray/experimental/channel/driver_communicator_holder.py
user: login: hipudding
body: Done.
created_at: 2025-03-19T07:01:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,364 @@
+import ctypes
+
+import torch
+
+hccl_lib = ctypes.CDLL("libhccl.so")
+
+HCCL_ROOT_INFO_BYTES = 4108
+
+
+class HcclRootInfo(ctypes.Structure):
+    _fields_ = [("internal", ctypes.c_byte * HCCL_ROOT_INFO_BYTES)]
+
+
+class HcclComm(ctypes.Structure):
+    pass
+
+
+class aclrtStream(ctypes.Structure):
path: python/ray/experimental/channel/hccl.py
user: login: hipudding
body: Done.
created_at: 2025-03-19T07:01:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -98,19 +98,212 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
 def get_default_torch_device(*, allow_cpu: bool) -> "torch.device":
     """Get the default torch device inside this actor or driver.
 
-    If any GPUs are available, the default device will be cuda:0 and we will rely on
-    torch to handle mapping CUDA_VISIBLE_DEVICES to a physical device.
+    If any acclerators are available, the default device will be index 0 and we
+    will rely on torch to handle mapping the VISIBLE_DEVICES to a physical device.
 
-    If no GPUs are available, a CPU device will be returned if allow_cpu is true, else
-    the function will raise a RuntimeError.
+    If no acclerators are available, a CPU device will be returned if allow_cpu
+    is true, else the function will raise a RuntimeError.
     """
     import torch
 
     accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
-    if not accelerator_ids.get("GPU", []):
+    if accelerator_ids.get("GPU", []):
+        return torch.device("cuda:0")
+    elif accelerator_ids.get("NPU", []):
+        return torch.device("npu:0")
+    else:
         if allow_cpu:
             return torch.device("cpu")
         else:
-            raise RuntimeError("No CUDA device available.")
+            raise RuntimeError("No acclerator available.")
+
+
+class AccleratorRuntime:
path: python/ray/experimental/channel/utils.py
user: login: noemotiovon
body: Maybe AcceleratorRuntime, other similar names are also recommended to be modified
created_at: 2025-03-20T01:38:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -98,19 +98,212 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
 def get_default_torch_device(*, allow_cpu: bool) -> "torch.device":
     """Get the default torch device inside this actor or driver.
 
-    If any GPUs are available, the default device will be cuda:0 and we will rely on
-    torch to handle mapping CUDA_VISIBLE_DEVICES to a physical device.
+    If any acclerators are available, the default device will be index 0 and we
+    will rely on torch to handle mapping the VISIBLE_DEVICES to a physical device.
 
-    If no GPUs are available, a CPU device will be returned if allow_cpu is true, else
-    the function will raise a RuntimeError.
+    If no acclerators are available, a CPU device will be returned if allow_cpu
+    is true, else the function will raise a RuntimeError.
     """
     import torch
 
     accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
-    if not accelerator_ids.get("GPU", []):
+    if accelerator_ids.get("GPU", []):
+        return torch.device("cuda:0")
+    elif accelerator_ids.get("NPU", []):
+        return torch.device("npu:0")
+    else:
         if allow_cpu:
             return torch.device("cpu")
         else:
-            raise RuntimeError("No CUDA device available.")
+            raise RuntimeError("No acclerator available.")
+
+
+class AccleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
+    """
+
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, NpuRuntime, or CpuRuntime).
+        """
+        if cls._instance is None:
+            device = get_default_torch_device(allow_cpu=True)
+            if device.type == "cuda":
+                cls._instance = CudaRuntime()
+            elif device.type == "npu":
+                cls._instance = NpuRuntime()
+            elif device.type == "cpu":
+                cls._instance = CpuRuntime()
+            else:
+                raise RuntimeError(f"Device type {device.type} is not support.")
+
+        return cls._instance
+
+    @classmethod
+    def is_acclerator_communicator_available(
+        self, device: Optional[Union["torch.device", str]] = None
+    ):
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device supports CUDA or NPU accelerators, False otherwise.
+        """
+        if device is None:
+            device = get_default_torch_device(allow_cpu=False)
+        if device in ["cuda", "npu"]:
+            return True
+        if device.type in ["cuda", "npu"]:
+            return True
+        return False
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device. Currently, it supports CUDA
+        and NPU devices.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
+        """
+        Generates a unique identifier for communication purposes.
+        """
+        raise NotImplementedError
+
+    def get_communicator(self, *args, **kwargs):
+        """
+        Retrieves the communication group for collective operations.
+        """
+        raise NotImplementedError
+
+
+class CudaRuntime(AccleratorRuntime):
+    """
+    CUDA implementation of AcceleratorRuntime.
+    """
+
+    def is_available(self) -> bool:
+        import torch
path: python/ray/experimental/channel/utils.py
user: login: noemotiovon
body: Multiple places require import torch. Would it be possible to import it within the current file?
created_at: 2025-03-20T01:50:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -98,19 +98,212 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
 def get_default_torch_device(*, allow_cpu: bool) -> "torch.device":
     """Get the default torch device inside this actor or driver.
 
-    If any GPUs are available, the default device will be cuda:0 and we will rely on
-    torch to handle mapping CUDA_VISIBLE_DEVICES to a physical device.
+    If any acclerators are available, the default device will be index 0 and we
+    will rely on torch to handle mapping the VISIBLE_DEVICES to a physical device.
 
-    If no GPUs are available, a CPU device will be returned if allow_cpu is true, else
-    the function will raise a RuntimeError.
+    If no acclerators are available, a CPU device will be returned if allow_cpu
+    is true, else the function will raise a RuntimeError.
     """
     import torch
 
     accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
-    if not accelerator_ids.get("GPU", []):
+    if accelerator_ids.get("GPU", []):
+        return torch.device("cuda:0")
+    elif accelerator_ids.get("NPU", []):
+        return torch.device("npu:0")
+    else:
         if allow_cpu:
             return torch.device("cpu")
         else:
-            raise RuntimeError("No CUDA device available.")
+            raise RuntimeError("No acclerator available.")
+
+
+class AccleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
+    """
+
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, NpuRuntime, or CpuRuntime).
+        """
+        if cls._instance is None:
+            device = get_default_torch_device(allow_cpu=True)
+            if device.type == "cuda":
+                cls._instance = CudaRuntime()
+            elif device.type == "npu":
+                cls._instance = NpuRuntime()
+            elif device.type == "cpu":
+                cls._instance = CpuRuntime()
+            else:
+                raise RuntimeError(f"Device type {device.type} is not support.")
+
+        return cls._instance
+
+    @classmethod
+    def is_acclerator_communicator_available(
+        self, device: Optional[Union["torch.device", str]] = None
path: python/ray/experimental/channel/utils.py
user: login: noemotiovon
body: cls, device: Optional[Union["torch.device", str]] = None
created_at: 2025-03-20T03:17:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -675,14 +675,16 @@ def _do_destroy_communicator(self, group_id):
     # task loop running.
 
 
-def _do_check_has_gpu(self) -> bool:
-    return bool(ray.get_gpu_ids())
+def _do_check_has_acclerator(self) -> bool:
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: noemotiovon
body: I feel that these methods are redundant. Would it be possible to directly import and call the AcceleratorRuntime-related methods where they are used?
created_at: 2025-03-20T03:44:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -191,7 +191,7 @@ def with_tensor_transport(
         else:
             if not isinstance(transport, Communicator):
                 raise ValueError(
-                    "transport must be 'auto', 'nccl' or a Communicator type"
+                    "transport must be 'auto', 'nccl', 'cpu' or a Communicator type"
path: python/ray/dag/dag_node.py
user: login: hipudding
body: This report is not the same with L176, which is the device check, this is transport check. I add cpu as a new transport since TorchTensorType support cpu.
created_at: 2025-03-20T12:00:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -98,19 +98,212 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
 def get_default_torch_device(*, allow_cpu: bool) -> "torch.device":
     """Get the default torch device inside this actor or driver.
 
-    If any GPUs are available, the default device will be cuda:0 and we will rely on
-    torch to handle mapping CUDA_VISIBLE_DEVICES to a physical device.
+    If any acclerators are available, the default device will be index 0 and we
+    will rely on torch to handle mapping the VISIBLE_DEVICES to a physical device.
 
-    If no GPUs are available, a CPU device will be returned if allow_cpu is true, else
-    the function will raise a RuntimeError.
+    If no acclerators are available, a CPU device will be returned if allow_cpu
+    is true, else the function will raise a RuntimeError.
     """
     import torch
 
     accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
-    if not accelerator_ids.get("GPU", []):
+    if accelerator_ids.get("GPU", []):
+        return torch.device("cuda:0")
+    elif accelerator_ids.get("NPU", []):
+        return torch.device("npu:0")
+    else:
         if allow_cpu:
             return torch.device("cpu")
         else:
-            raise RuntimeError("No CUDA device available.")
+            raise RuntimeError("No acclerator available.")
+
+
+class AccleratorRuntime:
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: Done.
created_at: 2025-03-25T06:37:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -98,19 +98,212 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
 def get_default_torch_device(*, allow_cpu: bool) -> "torch.device":
     """Get the default torch device inside this actor or driver.
 
-    If any GPUs are available, the default device will be cuda:0 and we will rely on
-    torch to handle mapping CUDA_VISIBLE_DEVICES to a physical device.
+    If any acclerators are available, the default device will be index 0 and we
+    will rely on torch to handle mapping the VISIBLE_DEVICES to a physical device.
 
-    If no GPUs are available, a CPU device will be returned if allow_cpu is true, else
-    the function will raise a RuntimeError.
+    If no acclerators are available, a CPU device will be returned if allow_cpu
+    is true, else the function will raise a RuntimeError.
     """
     import torch
 
     accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
-    if not accelerator_ids.get("GPU", []):
+    if accelerator_ids.get("GPU", []):
+        return torch.device("cuda:0")
+    elif accelerator_ids.get("NPU", []):
+        return torch.device("npu:0")
+    else:
         if allow_cpu:
             return torch.device("cpu")
         else:
-            raise RuntimeError("No CUDA device available.")
+            raise RuntimeError("No acclerator available.")
+
+
+class AccleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
+    """
+
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, NpuRuntime, or CpuRuntime).
+        """
+        if cls._instance is None:
+            device = get_default_torch_device(allow_cpu=True)
+            if device.type == "cuda":
+                cls._instance = CudaRuntime()
+            elif device.type == "npu":
+                cls._instance = NpuRuntime()
+            elif device.type == "cpu":
+                cls._instance = CpuRuntime()
+            else:
+                raise RuntimeError(f"Device type {device.type} is not support.")
+
+        return cls._instance
+
+    @classmethod
+    def is_acclerator_communicator_available(
+        self, device: Optional[Union["torch.device", str]] = None
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: Done.
created_at: 2025-03-25T06:37:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -98,19 +98,212 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
 def get_default_torch_device(*, allow_cpu: bool) -> "torch.device":
     """Get the default torch device inside this actor or driver.
 
-    If any GPUs are available, the default device will be cuda:0 and we will rely on
-    torch to handle mapping CUDA_VISIBLE_DEVICES to a physical device.
+    If any acclerators are available, the default device will be index 0 and we
+    will rely on torch to handle mapping the VISIBLE_DEVICES to a physical device.
 
-    If no GPUs are available, a CPU device will be returned if allow_cpu is true, else
-    the function will raise a RuntimeError.
+    If no acclerators are available, a CPU device will be returned if allow_cpu
+    is true, else the function will raise a RuntimeError.
     """
     import torch
 
     accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
-    if not accelerator_ids.get("GPU", []):
+    if accelerator_ids.get("GPU", []):
+        return torch.device("cuda:0")
+    elif accelerator_ids.get("NPU", []):
+        return torch.device("npu:0")
+    else:
         if allow_cpu:
             return torch.device("cpu")
         else:
-            raise RuntimeError("No CUDA device available.")
+            raise RuntimeError("No acclerator available.")
+
+
+class AccleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
+    """
+
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, NpuRuntime, or CpuRuntime).
+        """
+        if cls._instance is None:
+            device = get_default_torch_device(allow_cpu=True)
+            if device.type == "cuda":
+                cls._instance = CudaRuntime()
+            elif device.type == "npu":
+                cls._instance = NpuRuntime()
+            elif device.type == "cpu":
+                cls._instance = CpuRuntime()
+            else:
+                raise RuntimeError(f"Device type {device.type} is not support.")
+
+        return cls._instance
+
+    @classmethod
+    def is_acclerator_communicator_available(
+        self, device: Optional[Union["torch.device", str]] = None
+    ):
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device supports CUDA or NPU accelerators, False otherwise.
+        """
+        if device is None:
+            device = get_default_torch_device(allow_cpu=False)
+        if device in ["cuda", "npu"]:
+            return True
+        if device.type in ["cuda", "npu"]:
+            return True
+        return False
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device. Currently, it supports CUDA
+        and NPU devices.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
+        """
+        Generates a unique identifier for communication purposes.
+        """
+        raise NotImplementedError
+
+    def get_communicator(self, *args, **kwargs):
+        """
+        Retrieves the communication group for collective operations.
+        """
+        raise NotImplementedError
+
+
+class CudaRuntime(AccleratorRuntime):
+    """
+    CUDA implementation of AcceleratorRuntime.
+    """
+
+    def is_available(self) -> bool:
+        import torch
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: No, because Torch is not a dependency of Ray and is only referenced when using tensor collective communication.
created_at: 2025-03-25T08:08:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -675,14 +675,16 @@ def _do_destroy_communicator(self, group_id):
     # task loop running.
 
 
-def _do_check_has_gpu(self) -> bool:
-    return bool(ray.get_gpu_ids())
+def _do_check_has_acclerator(self) -> bool:
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: hipudding
body: This function will be used for remote calls by an Actor, so it’s better to keep it.
created_at: 2025-03-25T08:09:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -98,19 +98,164 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
 def get_default_torch_device(*, allow_cpu: bool) -> "torch.device":
     """Get the default torch device inside this actor or driver.
 
-    If any GPUs are available, the default device will be cuda:0 and we will rely on
-    torch to handle mapping CUDA_VISIBLE_DEVICES to a physical device.
+    If any acclerators are available, the default device will be index 0 and we
+    will rely on torch to handle mapping the VISIBLE_DEVICES to a physical device.
 
-    If no GPUs are available, a CPU device will be returned if allow_cpu is true, else
-    the function will raise a RuntimeError.
+    If no acclerators are available, a CPU device will be returned if allow_cpu
+    is true, else the function will raise a RuntimeError.
     """
     import torch
 
     accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
-    if not accelerator_ids.get("GPU", []):
+    if accelerator_ids.get("GPU", []):
+        return torch.device("cuda:0")
+    else:
         if allow_cpu:
             return torch.device("cpu")
         else:
-            raise RuntimeError("No CUDA device available.")
+            raise RuntimeError("No acclerator available.")
+
+
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
+    """
+
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, NpuRuntime, or CpuRuntime).
+        """
+        if cls._instance is None:
+            device = get_default_torch_device(allow_cpu=True)
+            if device.type == "cuda":
+                cls._instance = CudaRuntime()
+            elif device.type == "cpu":
+                cls._instance = CpuRuntime()
+            else:
+                raise RuntimeError(f"Device type {device.type} is not support.")
+
+        return cls._instance
+
+    @classmethod
+    def has_communicator(cls, device: Optional[Union["torch.device", str]] = None):
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device supports CUDA or NPU accelerators, False otherwise.
+        """
+        if device is None:
+            device = get_default_torch_device(allow_cpu=False)
+        if not isinstance(device, str):
+            device = device.type
+        return device in ["cuda"]
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device. Currently, it supports CUDA
+        and NPU devices.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
+        """
+        Generates a unique identifier for communication purposes.
+        """
+        raise NotImplementedError
+
+    def get_communicator(self, *args, **kwargs):
+        """
+        Retrieves the communication group for collective operations.
+        """
+        raise NotImplementedError
+
+
+class CudaRuntime(AcceleratorRuntime):
+    """
+    CUDA implementation of AcceleratorRuntime.
+    """
+
+    def is_available(self) -> bool:
+        import torch
+
+        return torch.cuda.is_available()
+
+    def current_stream(self):
+        import torch
+
+        return torch.cuda.current_stream()
+
+    def create_event(self):
+        import torch
+
+        return torch.cuda.Event()
+
+    def get_device_context(self, device):
+        import torch
+
+        return torch.cuda.device(device)
+
+    def get_unique_id(self) -> str:
+        from ray.experimental.channel.nccl_group import get_nccl_unique_id
+
+        return get_nccl_unique_id()
+
+    def get_communicator(self, *args, **kwargs):
+        from ray.experimental.channel.nccl_group import _NcclGroup
+
+        return _NcclGroup(*args, **kwargs)
+
+
+class CpuRuntime(AcceleratorRuntime):
+    """
+    CPU implementation of AcceleratorRuntime.
+    """
+
+    def get_unique_id(self) -> str:
+        from ray.experimental.channel.cpu_communicator import generate_uid_id
 
-    return torch.device("cuda:0")
+        return generate_uid_id()
path: python/ray/experimental/channel/utils.py
user: login: noemotiovon
body: Could this method be directly inlined? Or alternatively, could the `get_unique_id` method in the base class be defined as str(uuid.uuid4())
created_at: 2025-03-27T02:02:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,91 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorInfoHolder(Communicator):
+    """
+    The driver maintains a communicator instance but does not participate
+    in the group. However, it still needs to store certain information,
+    such as world size and actor details.
+
+    Since the driver may not have a GPU (or other accelerators), it could
+    mistakenly select the wrong `TorchDeviceManager`
+    (e.g., `CPUTorchDeviceManager`), causing `get_communicator`
+    to return `None`.
+
+    Because the driver does not actively join the group and only uses the
+    communicator for metadata storage, Introduced a hardware-independent
+    class specifically for the driver.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        comm_id: int,
path: python/ray/experimental/channel/driver_communicator_holder.py
user: login: noemotiovon
body: This parameter is unused.
created_at: 2025-03-27T02:11:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,91 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorInfoHolder(Communicator):
+    """
+    The driver maintains a communicator instance but does not participate
+    in the group. However, it still needs to store certain information,
+    such as world size and actor details.
+
+    Since the driver may not have a GPU (or other accelerators), it could
+    mistakenly select the wrong `TorchDeviceManager`
+    (e.g., `CPUTorchDeviceManager`), causing `get_communicator`
+    to return `None`.
+
+    Because the driver does not actively join the group and only uses the
+    communicator for metadata storage, Introduced a hardware-independent
+    class specifically for the driver.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        comm_id: int,
path: python/ray/experimental/channel/driver_communicator_holder.py
user: login: hipudding
body: Done.
created_at: 2025-03-28T01:31:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -98,19 +98,164 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
 def get_default_torch_device(*, allow_cpu: bool) -> "torch.device":
     """Get the default torch device inside this actor or driver.
 
-    If any GPUs are available, the default device will be cuda:0 and we will rely on
-    torch to handle mapping CUDA_VISIBLE_DEVICES to a physical device.
+    If any acclerators are available, the default device will be index 0 and we
+    will rely on torch to handle mapping the VISIBLE_DEVICES to a physical device.
 
-    If no GPUs are available, a CPU device will be returned if allow_cpu is true, else
-    the function will raise a RuntimeError.
+    If no acclerators are available, a CPU device will be returned if allow_cpu
+    is true, else the function will raise a RuntimeError.
     """
     import torch
 
     accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
-    if not accelerator_ids.get("GPU", []):
+    if accelerator_ids.get("GPU", []):
+        return torch.device("cuda:0")
+    else:
         if allow_cpu:
             return torch.device("cpu")
         else:
-            raise RuntimeError("No CUDA device available.")
+            raise RuntimeError("No acclerator available.")
+
+
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
+    """
+
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, NpuRuntime, or CpuRuntime).
+        """
+        if cls._instance is None:
+            device = get_default_torch_device(allow_cpu=True)
+            if device.type == "cuda":
+                cls._instance = CudaRuntime()
+            elif device.type == "cpu":
+                cls._instance = CpuRuntime()
+            else:
+                raise RuntimeError(f"Device type {device.type} is not support.")
+
+        return cls._instance
+
+    @classmethod
+    def has_communicator(cls, device: Optional[Union["torch.device", str]] = None):
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device supports CUDA or NPU accelerators, False otherwise.
+        """
+        if device is None:
+            device = get_default_torch_device(allow_cpu=False)
+        if not isinstance(device, str):
+            device = device.type
+        return device in ["cuda"]
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device. Currently, it supports CUDA
+        and NPU devices.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
+        """
+        Generates a unique identifier for communication purposes.
+        """
+        raise NotImplementedError
+
+    def get_communicator(self, *args, **kwargs):
+        """
+        Retrieves the communication group for collective operations.
+        """
+        raise NotImplementedError
+
+
+class CudaRuntime(AcceleratorRuntime):
+    """
+    CUDA implementation of AcceleratorRuntime.
+    """
+
+    def is_available(self) -> bool:
+        import torch
+
+        return torch.cuda.is_available()
+
+    def current_stream(self):
+        import torch
+
+        return torch.cuda.current_stream()
+
+    def create_event(self):
+        import torch
+
+        return torch.cuda.Event()
+
+    def get_device_context(self, device):
+        import torch
+
+        return torch.cuda.device(device)
+
+    def get_unique_id(self) -> str:
+        from ray.experimental.channel.nccl_group import get_nccl_unique_id
+
+        return get_nccl_unique_id()
+
+    def get_communicator(self, *args, **kwargs):
+        from ray.experimental.channel.nccl_group import _NcclGroup
+
+        return _NcclGroup(*args, **kwargs)
+
+
+class CpuRuntime(AcceleratorRuntime):
+    """
+    CPU implementation of AcceleratorRuntime.
+    """
+
+    def get_unique_id(self) -> str:
+        from ray.experimental.channel.cpu_communicator import generate_uid_id
 
-    return torch.device("cuda:0")
+        return generate_uid_id()
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: Fixed.
created_at: 2025-03-28T01:31:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
path: python/ray/experimental/channel/utils.py
user: login: ruisearch42
body: These abstractions should be put into a different file rather than `utils.py`
created_at: 2025-03-31T18:20:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
path: python/ray/experimental/channel/utils.py
user: login: ruisearch42
body: This is not thread-safe. Need to think a bit more carefully if we need locking, or just add a comment.
created_at: 2025-03-31T18:21:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
path: python/ray/experimental/channel/utils.py
user: login: ruisearch42
body: arg does not match function signature
created_at: 2025-03-31T18:28:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
path: python/ray/experimental/channel/utils.py
user: login: ruisearch42
body: Can you please rephrase this to be clearer?
created_at: 2025-03-31T18:29:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
path: python/ray/experimental/channel/utils.py
user: login: ruisearch42
body: Can we have a type hint for the return value?
created_at: 2025-03-31T18:31:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
+        """
+        Generates a unique identifier for communication purposes.
path: python/ray/experimental/channel/utils.py
user: login: ruisearch42
body: Why not call it `generate_unique_id` or `create_unique_id`?
created_at: 2025-03-31T18:31:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
+        """
+        Generates a unique identifier for communication purposes.
+        """
+        raise NotImplementedError
+
+    def get_communicator(self, *args, **kwargs):
+        """
+        Retrieves the communication group for collective operations.
+        """
+        raise NotImplementedError
+
+    def get_device_type(self) -> str:
+        """
+        Retrieves the device type for the accelerator.
path: python/ray/experimental/channel/utils.py
user: login: ruisearch42
body: If this is device type from PyTorch, mention that
created_at: 2025-03-31T18:32:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,90 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorInfoHolder(Communicator):
+    """
+    The driver maintains a communicator instance but does not participate
+    in the group. However, it still needs to store certain information,
+    such as world size and actor details.
+
+    Since the driver may not have a GPU (or other accelerators), it could
+    mistakenly select the wrong `TorchDeviceManager`
+    (e.g., `CPUTorchDeviceManager`), causing `get_communicator`
+    to return `None`.
+
+    Because the driver does not actively join the group and only uses the
+    communicator for metadata storage, Introduced a hardware-independent
+    class specifically for the driver.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        actor_handles: List["ray.actor.ActorHandle"],
+    ):
+        """
+        Only need to store the actor handles, since we don't need
+        to initialize the communicator for Driver.
+        """
+        self._world_size = world_size
+        self._actor_handles = actor_handles
+
+    def initialize(self, rank: int) -> None:
+        # No additional initialization is needed.
+        pass
+
+    def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        """
+        Retuan all actor handles in this communicator.
+        """
+        return self._actor_handles
+
+    def get_rank(self, actor: ray.actor.ActorHandle) -> int:
+        return None
+
+    def get_self_rank(self) -> Optional[int]:
+        return None
+
+    def send(self, tensor: "torch.Tensor", peer_rank: int):
+        raise NotImplementedError
+
+    def recv(
+        self,
+        shape: Tuple[int],
+        dtype: "torch.dtype",
+        peer_rank: int,
+        allocator: Optional[TorchTensorAllocator] = None,
+    ):
+        raise NotImplementedError
+
+    def allreduce(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp = ReduceOp.SUM,
+    ):
+        raise NotImplementedError
+
+    def destroy(self) -> None:
+        pass
+
+    def get_world_size(self) -> int:
+        """
+        Return the number of ranks in the communicator.
+        """
+        return self._world_size
+
+    def get_transport_name(self) -> str:
+        raise NotImplementedError
+
+    def recv_stream(self):
+        raise NotImplementedError
+
+    def send_stream(self):
+        raise NotImplementedError
path: python/ray/experimental/channel/communicator_info_holder.py
user: login: ruisearch42
body: you don't need these, right?
Also for other methods that already have `raise NotImplementedError` in the base class
created_at: 2025-03-31T18:40:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,90 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorInfoHolder(Communicator):
path: python/ray/experimental/channel/communicator_info_holder.py
user: login: ruisearch42
body: Can we improve the naming a bit? I'm thinking about `CommunicatorMetadata`, but it's still not accurate enough
created_at: 2025-03-31T22:25:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
path: python/ray/experimental/channel/utils.py
user: login: ruisearch42
body: How about `AcceleratorContext`?
created_at: 2025-03-31T22:29:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -591,8 +590,8 @@ def prepare(self, overlap_gpu_communication: bool = False):
         self.input_reader.start()
         self.output_writer.start()
 
-        self._send_stream: Union["cp.cuda.Stream", nullcontext] = nullcontext()
-        self._recv_stream: Union["cp.cuda.Stream", nullcontext] = nullcontext()
path: python/ray/dag/compiled_dag_node.py
user: login: stephanie-wang
body: Let's still keep a generic type to indicate that this is a context?
created_at: 2025-04-06T22:03:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -100,25 +97,23 @@ def remove_gpu_future(fut_id: int) -> None:
         if fut_id in GPUFuture.gpu_futures:
             GPUFuture.gpu_futures.pop(fut_id).destroy_event()
 
-    def __init__(
-        self, buf: Any, fut_id: int, stream: Optional["cp.cuda.Stream"] = None
-    ):
+    def __init__(self, buf: Any, fut_id: int, stream: Any = None):
         """
         Initialize a GPU future on the given stream.
 
         Args:
             buf: The buffer to return when the future is resolved.
             fut_id: The future ID to cache the future.
-            stream: The CUDA stream to record the event on, this event is waited
+            stream: The torch stream to record the event on, this event is waited
                 on when the future is resolved. If None, the current stream is used.
         """
-        import cupy as cp
+        from ray.experimental.channel.utils import AcceleratorRuntime
path: python/ray/dag/dag_operation_future.py
user: login: stephanie-wang
body: Let's import at top-level. Importing `cupy` here was only necessary because we don't want a hard dependency on it.
created_at: 2025-04-06T22:03:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -146,11 +141,7 @@ def destroy_event(self) -> None:
         """
         Destroy the CUDA event associated with this future.
         """
-        import cupy as cp
-
         if self._event is None:
             return
 
-        cp.cuda.runtime.eventDestroy(self._event.ptr)
path: python/ray/dag/dag_operation_future.py
user: login: stephanie-wang
body: Can you bring this code back? Without this, the CUDA event can get destroyed before the stream and it may segfault.
created_at: 2025-04-06T22:04:30Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
path: python/ray/experimental/channel/utils.py
user: login: stephanie-wang
body: I am a bit confused about what this method is supposed to do. Why does it always return True for CUDA and False for CPU?
created_at: 2025-04-06T22:06:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
path: python/ray/experimental/channel/utils.py
user: login: stephanie-wang
body: It is a little confusing that the runtime can be created but not available. I guess this is meant to return whether the device is available for individual actors? Does it ever return anything different from `len(get_devices()) > 0`?
created_at: 2025-04-06T22:07:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
+        """
+        Generates a unique identifier for communication purposes.
+        """
+        raise NotImplementedError
+
+    def get_communicator(self, *args, **kwargs):
+        """
+        Retrieves the communication group for collective operations.
path: python/ray/experimental/channel/utils.py
user: login: stephanie-wang
body: What if the actor is in multiple communicator groups?
created_at: 2025-04-06T22:07:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
path: python/ray/experimental/channel/utils.py
user: login: stephanie-wang
body: ```suggestion
    def generate_communicator_id(self) -> str:
```
created_at: 2025-04-06T22:09:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
path: python/ray/experimental/channel/utils.py
user: login: stephanie-wang
body: (Let's put communicator in the name)
created_at: 2025-04-06T22:09:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,90 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorInfoHolder(Communicator):
path: python/ray/experimental/channel/communicator_info_holder.py
user: login: stephanie-wang
body: Yes, let's call this something different. How about CommunicatorHandle to match ActorHandle?
created_at: 2025-04-06T22:10:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,90 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorInfoHolder(Communicator):
+    """
+    The driver maintains a communicator instance but does not participate
+    in the group. However, it still needs to store certain information,
+    such as world size and actor details.
+
+    Since the driver may not have a GPU (or other accelerators), it could
+    mistakenly select the wrong `TorchDeviceManager`
+    (e.g., `CPUTorchDeviceManager`), causing `get_communicator`
+    to return `None`.
+
+    Because the driver does not actively join the group and only uses the
+    communicator for metadata storage, Introduced a hardware-independent
+    class specifically for the driver.
path: python/ray/experimental/channel/communicator_info_holder.py
user: login: stephanie-wang
body: This comment is written in a bit of a confusing way because it is trying to justify why we need this class, so it only makes sense in the context of this PR.

Instead, can you just write what this class is responsible for?
created_at: 2025-04-06T22:12:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,90 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorInfoHolder(Communicator):
+    """
+    The driver maintains a communicator instance but does not participate
+    in the group. However, it still needs to store certain information,
+    such as world size and actor details.
+
+    Since the driver may not have a GPU (or other accelerators), it could
+    mistakenly select the wrong `TorchDeviceManager`
+    (e.g., `CPUTorchDeviceManager`), causing `get_communicator`
+    to return `None`.
+
+    Because the driver does not actively join the group and only uses the
+    communicator for metadata storage, Introduced a hardware-independent
+    class specifically for the driver.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        actor_handles: List["ray.actor.ActorHandle"],
+    ):
+        """
+        Only need to store the actor handles, since we don't need
+        to initialize the communicator for Driver.
+        """
+        self._world_size = world_size
+        self._actor_handles = actor_handles
+
+    def initialize(self, rank: int) -> None:
+        # No additional initialization is needed.
+        pass
+
+    def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        """
+        Retuan all actor handles in this communicator.
+        """
+        return self._actor_handles
+
+    def get_rank(self, actor: ray.actor.ActorHandle) -> int:
+        return None
+
+    def get_self_rank(self) -> Optional[int]:
+        return None
+
+    def send(self, tensor: "torch.Tensor", peer_rank: int):
path: python/ray/experimental/channel/communicator_info_holder.py
user: login: stephanie-wang
body: Should we remove these? They seem unnecessary from the driver's perspective?
created_at: 2025-04-06T22:13:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
+        """
+        Generates a unique identifier for communication purposes.
+        """
+        raise NotImplementedError
+
+    def get_communicator(self, *args, **kwargs):
+        """
+        Retrieves the communication group for collective operations.
+        """
+        raise NotImplementedError
+
+    def get_device_type(self) -> str:
+        """
+        Retrieves the device type for the accelerator.
+        """
+        raise NotImplementedError
+
+
+class CudaRuntime(AcceleratorRuntime):
+    """
+    CUDA implementation of AcceleratorRuntime.
+    """
 
-    device_ids = []
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch cuda device list configured for this process.
+
+        Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
+        superset of the `ray.get_gpu_ids()`.
+        """
+        # Note: currently this method replicates the logic from
+        # `CUDATorchDeviceManager.get_devices()`.
+        # TODO(rui): tailor and clean up the logic for proper use in
+        # Compiled Graphs.
+        import torch
+
+        # GPU IDs are assigned by Ray after you specify "use_gpu"
+        # GPU `ray.get_gpu_ids()` may return ints or may return strings.
+        # We should always convert to strings.
+        gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+
+        device_ids = []
+
+        if len(gpu_ids) > 0:
+            cuda_visible_str = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+            if cuda_visible_str and cuda_visible_str != "NoDevFiles":
+                cuda_visible_list = cuda_visible_str.split(",")
+            else:
+                cuda_visible_list = []
+
+            # By default, there should only be one GPU ID if `use_gpu=True`.
+            # If there are multiple GPUs, return a list of devices.
+            # If using fractional GPUs, these IDs are not guaranteed
+            # to be unique across different processes.
+            for gpu_id in gpu_ids:
+                try:
+                    device_ids.append(cuda_visible_list.index(gpu_id))
+                except IndexError:
+                    raise RuntimeError(
+                        "CUDA_VISIBLE_DEVICES set incorrectly. "
+                        f"Got {cuda_visible_str}, expected to include {gpu_id}. "
+                        "Did you override the `CUDA_VISIBLE_DEVICES` environment"
+                        " variable? If not, please help file an issue on Github."
+                    )
 
-    if len(gpu_ids) > 0:
-        cuda_visible_str = os.environ.get("CUDA_VISIBLE_DEVICES", "")
-        if cuda_visible_str and cuda_visible_str != "NoDevFiles":
-            cuda_visible_list = cuda_visible_str.split(",")
         else:
-            cuda_visible_list = []
-
-        # By default, there should only be one GPU ID if `use_gpu=True`.
-        # If there are multiple GPUs, return a list of devices.
-        # If using fractional GPUs, these IDs are not guaranteed
-        # to be unique across different processes.
-        for gpu_id in gpu_ids:
-            try:
-                device_ids.append(cuda_visible_list.index(gpu_id))
-            except IndexError:
-                raise RuntimeError(
-                    "CUDA_VISIBLE_DEVICES set incorrectly. "
-                    f"Got {cuda_visible_str}, expected to include {gpu_id}. "
-                    "Did you override the `CUDA_VISIBLE_DEVICES` environment"
-                    " variable? If not, please help file an issue on Github."
-                )
+            # If called on the driver or outside of Ray Train, return the
+            # 0th device.
+            device_ids.append(0)
 
-    else:
-        # If called on the driver or outside of Ray Train, return the
-        # 0th device.
-        device_ids.append(0)
+        return [torch.device(f"cuda:{device_id}") for device_id in device_ids]
+
+    def has_communicator(self):
+        return True
+
+    def is_available(self) -> bool:
+        import torch
+
+        return torch.cuda.is_available()
+
+    def current_stream(self):
+        import torch
+
+        return torch.cuda.current_stream()
+
+    def create_event(self):
+        import torch
+
+        return torch.cuda.Event()
+
+    def get_device_context(self, device):
+        import torch
 
-    return [torch.device(f"cuda:{device_id}") for device_id in device_ids]
+        return torch.cuda.device(device)
 
+    def get_unique_id(self) -> str:
+        from ray.experimental.channel.nccl_group import get_nccl_unique_id
 
-def get_devices() -> List["torch.device"]:
-    """Gets the correct torch device list configured for this process.
+        return get_nccl_unique_id()
 
-    Returns a list of torch devices allocated for the current worker.
-    If no devices are assigned, then it returns a list with a single CPU device.
+    def get_communicator(self, *args, **kwargs):
+        from ray.experimental.channel.nccl_group import _NcclGroup
+
+        return _NcclGroup(*args, **kwargs)
+
+    def get_device_type(self) -> str:
+        return "cuda"
+
+
+class CpuRuntime(AcceleratorRuntime):
path: python/ray/experimental/channel/utils.py
user: login: stephanie-wang
body: I am a bit confused how the CpuRuntime fits in? Is it meant to be used as a default runtime? What happens if an actor requests 1 CPU and 1 GPU, will it have both a CpuRuntime and a CudaRuntime?
created_at: 2025-04-06T22:14:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: Here, I don’t think it’s possible to add a return value type hint, because the type of the device varies across different backend devices — for example, torch.cuda.device and torch.npu.device.
created_at: 2025-04-22T08:28:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,90 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorInfoHolder(Communicator):
+    """
+    The driver maintains a communicator instance but does not participate
+    in the group. However, it still needs to store certain information,
+    such as world size and actor details.
+
+    Since the driver may not have a GPU (or other accelerators), it could
+    mistakenly select the wrong `TorchDeviceManager`
+    (e.g., `CPUTorchDeviceManager`), causing `get_communicator`
+    to return `None`.
+
+    Because the driver does not actively join the group and only uses the
+    communicator for metadata storage, Introduced a hardware-independent
+    class specifically for the driver.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        actor_handles: List["ray.actor.ActorHandle"],
+    ):
+        """
+        Only need to store the actor handles, since we don't need
+        to initialize the communicator for Driver.
+        """
+        self._world_size = world_size
+        self._actor_handles = actor_handles
+
+    def initialize(self, rank: int) -> None:
+        # No additional initialization is needed.
+        pass
+
+    def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        """
+        Retuan all actor handles in this communicator.
+        """
+        return self._actor_handles
+
+    def get_rank(self, actor: ray.actor.ActorHandle) -> int:
+        return None
+
+    def get_self_rank(self) -> Optional[int]:
+        return None
+
+    def send(self, tensor: "torch.Tensor", peer_rank: int):
+        raise NotImplementedError
+
+    def recv(
+        self,
+        shape: Tuple[int],
+        dtype: "torch.dtype",
+        peer_rank: int,
+        allocator: Optional[TorchTensorAllocator] = None,
+    ):
+        raise NotImplementedError
+
+    def allreduce(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp = ReduceOp.SUM,
+    ):
+        raise NotImplementedError
+
+    def destroy(self) -> None:
+        pass
+
+    def get_world_size(self) -> int:
+        """
+        Return the number of ranks in the communicator.
+        """
+        return self._world_size
+
+    def get_transport_name(self) -> str:
+        raise NotImplementedError
+
+    def recv_stream(self):
+        raise NotImplementedError
+
+    def send_stream(self):
+        raise NotImplementedError
path: python/ray/experimental/channel/communicator_info_holder.py
user: login: hipudding
body: No. I don't need this, but these function are marked as abstractmethod in base class. 
created_at: 2025-04-22T08:32:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -100,25 +97,23 @@ def remove_gpu_future(fut_id: int) -> None:
         if fut_id in GPUFuture.gpu_futures:
             GPUFuture.gpu_futures.pop(fut_id).destroy_event()
 
-    def __init__(
-        self, buf: Any, fut_id: int, stream: Optional["cp.cuda.Stream"] = None
-    ):
+    def __init__(self, buf: Any, fut_id: int, stream: Any = None):
         """
         Initialize a GPU future on the given stream.
 
         Args:
             buf: The buffer to return when the future is resolved.
             fut_id: The future ID to cache the future.
-            stream: The CUDA stream to record the event on, this event is waited
+            stream: The torch stream to record the event on, this event is waited
                 on when the future is resolved. If None, the current stream is used.
         """
-        import cupy as cp
+        from ray.experimental.channel.utils import AcceleratorRuntime
path: python/ray/dag/dag_operation_future.py
user: login: hipudding
body: For keep a generic type to indicate for _send_stream and _recv_stream:
I can’t specify a type hint here because different backends have different Stream Contexts — for example, torch.cuda.StreamContext or torch.npu.StreamContext.
created_at: 2025-04-22T08:37:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -146,11 +141,7 @@ def destroy_event(self) -> None:
         """
         Destroy the CUDA event associated with this future.
         """
-        import cupy as cp
-
         if self._event is None:
             return
 
-        cp.cuda.runtime.eventDestroy(self._event.ptr)
path: python/ray/dag/dag_operation_future.py
user: login: hipudding
body: I’ve read the relevant PR. I believe this won’t introduce any issues:
	1.	When self_event is set to None, it triggers the destruction of the object. In the destructor, eventDestroy is called. According to the documentation I found, this destruction process is performed synchronously.
	2.	This PR lets torch manage the stream, device, and event. Since torch manages a stream pool and streams are not destroyed, the timing of releasing the event here will not cause a segmentation fault.
created_at: 2025-04-22T08:42:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: This function is used to determine whether the current backend supports out-of-band communication. CUDA supports communication via NCCL, while the CPU always uses shared memory for data transfer.
created_at: 2025-04-22T08:47:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
+        """
+        Generates a unique identifier for communication purposes.
+        """
+        raise NotImplementedError
+
+    def get_communicator(self, *args, **kwargs):
+        """
+        Retrieves the communication group for collective operations.
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: This has no impact. get_communicator will create a communication group based on the given actor list. If an actor is part of multiple communicator groups, this function will be called multiple times.
created_at: 2025-04-22T08:52:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,90 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorInfoHolder(Communicator):
+    """
+    The driver maintains a communicator instance but does not participate
+    in the group. However, it still needs to store certain information,
+    such as world size and actor details.
+
+    Since the driver may not have a GPU (or other accelerators), it could
+    mistakenly select the wrong `TorchDeviceManager`
+    (e.g., `CPUTorchDeviceManager`), causing `get_communicator`
+    to return `None`.
+
+    Because the driver does not actively join the group and only uses the
+    communicator for metadata storage, Introduced a hardware-independent
+    class specifically for the driver.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        actor_handles: List["ray.actor.ActorHandle"],
+    ):
+        """
+        Only need to store the actor handles, since we don't need
+        to initialize the communicator for Driver.
+        """
+        self._world_size = world_size
+        self._actor_handles = actor_handles
+
+    def initialize(self, rank: int) -> None:
+        # No additional initialization is needed.
+        pass
+
+    def get_actor_handles(self) -> List["ray.actor.ActorHandle"]:
+        """
+        Retuan all actor handles in this communicator.
+        """
+        return self._actor_handles
+
+    def get_rank(self, actor: ray.actor.ActorHandle) -> int:
+        return None
+
+    def get_self_rank(self) -> Optional[int]:
+        return None
+
+    def send(self, tensor: "torch.Tensor", peer_rank: int):
path: python/ray/experimental/channel/communicator_info_holder.py
user: login: hipudding
body: I'd like to. But these functions are marked abstrctmethod.
created_at: 2025-04-22T08:53:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
+        """
+        Generates a unique identifier for communication purposes.
+        """
+        raise NotImplementedError
+
+    def get_communicator(self, *args, **kwargs):
+        """
+        Retrieves the communication group for collective operations.
+        """
+        raise NotImplementedError
+
+    def get_device_type(self) -> str:
+        """
+        Retrieves the device type for the accelerator.
+        """
+        raise NotImplementedError
+
+
+class CudaRuntime(AcceleratorRuntime):
+    """
+    CUDA implementation of AcceleratorRuntime.
+    """
 
-    device_ids = []
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch cuda device list configured for this process.
+
+        Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
+        superset of the `ray.get_gpu_ids()`.
+        """
+        # Note: currently this method replicates the logic from
+        # `CUDATorchDeviceManager.get_devices()`.
+        # TODO(rui): tailor and clean up the logic for proper use in
+        # Compiled Graphs.
+        import torch
+
+        # GPU IDs are assigned by Ray after you specify "use_gpu"
+        # GPU `ray.get_gpu_ids()` may return ints or may return strings.
+        # We should always convert to strings.
+        gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+
+        device_ids = []
+
+        if len(gpu_ids) > 0:
+            cuda_visible_str = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+            if cuda_visible_str and cuda_visible_str != "NoDevFiles":
+                cuda_visible_list = cuda_visible_str.split(",")
+            else:
+                cuda_visible_list = []
+
+            # By default, there should only be one GPU ID if `use_gpu=True`.
+            # If there are multiple GPUs, return a list of devices.
+            # If using fractional GPUs, these IDs are not guaranteed
+            # to be unique across different processes.
+            for gpu_id in gpu_ids:
+                try:
+                    device_ids.append(cuda_visible_list.index(gpu_id))
+                except IndexError:
+                    raise RuntimeError(
+                        "CUDA_VISIBLE_DEVICES set incorrectly. "
+                        f"Got {cuda_visible_str}, expected to include {gpu_id}. "
+                        "Did you override the `CUDA_VISIBLE_DEVICES` environment"
+                        " variable? If not, please help file an issue on Github."
+                    )
 
-    if len(gpu_ids) > 0:
-        cuda_visible_str = os.environ.get("CUDA_VISIBLE_DEVICES", "")
-        if cuda_visible_str and cuda_visible_str != "NoDevFiles":
-            cuda_visible_list = cuda_visible_str.split(",")
         else:
-            cuda_visible_list = []
-
-        # By default, there should only be one GPU ID if `use_gpu=True`.
-        # If there are multiple GPUs, return a list of devices.
-        # If using fractional GPUs, these IDs are not guaranteed
-        # to be unique across different processes.
-        for gpu_id in gpu_ids:
-            try:
-                device_ids.append(cuda_visible_list.index(gpu_id))
-            except IndexError:
-                raise RuntimeError(
-                    "CUDA_VISIBLE_DEVICES set incorrectly. "
-                    f"Got {cuda_visible_str}, expected to include {gpu_id}. "
-                    "Did you override the `CUDA_VISIBLE_DEVICES` environment"
-                    " variable? If not, please help file an issue on Github."
-                )
+            # If called on the driver or outside of Ray Train, return the
+            # 0th device.
+            device_ids.append(0)
 
-    else:
-        # If called on the driver or outside of Ray Train, return the
-        # 0th device.
-        device_ids.append(0)
+        return [torch.device(f"cuda:{device_id}") for device_id in device_ids]
+
+    def has_communicator(self):
+        return True
+
+    def is_available(self) -> bool:
+        import torch
+
+        return torch.cuda.is_available()
+
+    def current_stream(self):
+        import torch
+
+        return torch.cuda.current_stream()
+
+    def create_event(self):
+        import torch
+
+        return torch.cuda.Event()
+
+    def get_device_context(self, device):
+        import torch
 
-    return [torch.device(f"cuda:{device_id}") for device_id in device_ids]
+        return torch.cuda.device(device)
 
+    def get_unique_id(self) -> str:
+        from ray.experimental.channel.nccl_group import get_nccl_unique_id
 
-def get_devices() -> List["torch.device"]:
-    """Gets the correct torch device list configured for this process.
+        return get_nccl_unique_id()
 
-    Returns a list of torch devices allocated for the current worker.
-    If no devices are assigned, then it returns a list with a single CPU device.
+    def get_communicator(self, *args, **kwargs):
+        from ray.experimental.channel.nccl_group import _NcclGroup
+
+        return _NcclGroup(*args, **kwargs)
+
+    def get_device_type(self) -> str:
+        return "cuda"
+
+
+class CpuRuntime(AcceleratorRuntime):
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: Compiled graphs can still be used even without a GPU. In this case, torch.tensor will be transferred via shared memory. When AcceleratorRuntime is invoked for the first time, it will first check for the presence of a CUDA device. If none is found, it will default to using the CPU.
created_at: 2025-04-22T08:56:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: Done.
created_at: 2025-04-29T02:24:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: Done.
created_at: 2025-04-29T02:25:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: Done.
created_at: 2025-04-29T02:25:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: This function has beed removed.
created_at: 2025-04-29T02:25:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
+        """
+        Generates a unique identifier for communication purposes.
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: Done.
created_at: 2025-04-29T02:25:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
+        """
+        Generates a unique identifier for communication purposes.
+        """
+        raise NotImplementedError
+
+    def get_communicator(self, *args, **kwargs):
+        """
+        Retrieves the communication group for collective operations.
+        """
+        raise NotImplementedError
+
+    def get_device_type(self) -> str:
+        """
+        Retrieves the device type for the accelerator.
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: This function has beed removed.
created_at: 2025-04-29T02:26:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,90 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorInfoHolder(Communicator):
path: python/ray/experimental/channel/communicator_info_holder.py
user: login: hipudding
body: Done.
created_at: 2025-04-29T02:26:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: Done.
created_at: 2025-04-29T02:26:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: This function has beed removed.
created_at: 2025-04-29T02:26:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -96,66 +96,226 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
         )
 
 
-def get_cuda_devices() -> List["torch.device"]:
-    """Gets the correct torch cuda device list configured for this process.
-
-    Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
-    superset of the `ray.get_gpu_ids()`.
+class AcceleratorRuntime:
+    """
+    Base class for managing device-specific runtime operations. This class
+    provides an interface for device-related operations such as stream
+    management, event creation, and device communication.
     """
-    # Note: currently this method replicates the logic from
-    # `CUDATorchDeviceManager.get_devices()`.
-    # TODO(rui): tailor and clean up the logic for proper use in
-    # Compiled Graphs.
-    import torch
 
-    # GPU IDs are assigned by Ray after you specify "use_gpu"
-    # GPU `ray.get_gpu_ids()` may return ints or may return strings.
-    # We should always convert to strings.
-    gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+    _instance = None
+
+    @classmethod
+    def get(cls):
+        """
+        Singleton method to get the runtime instance based on the default torch
+        device. It initializes the appropriate runtime class based on the
+        detected device type.
+        Returns:
+            AcceleratorRuntime: An instance of the appropriate runtime class
+            (CudaRuntime, or CpuRuntime).
+        """
+
+        if cls._instance is None:
+            if len([str(id) for id in ray.get_gpu_ids()]) > 0:
+                cls._instance = CudaRuntime()
+            else:
+                cls._instance = CpuRuntime()
+
+        return cls._instance
+
+    def get_devices(self) -> List["torch.device"]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch devices allocated for the current worker.
+        If no devices are assigned, then it returns a list with a single CPU device.
+        """
+        raise NotImplementedError
+
+    def has_communicator(self) -> bool:
+        """
+        Check if an accelerator communicator is available for the specified device.
+        If no device is provided, the function determines the default non-CPU device.
+        Args:
+            device (Optional[Union["torch.device", str]]):
+                The target device or device type, which can be a `torch.device`
+                object or a string. If not provided, the function uses the default
+                non-CPU device.
+        Returns:
+            bool: True if the device support out of band communication, False otherwise.
+        """
+        raise NotImplementedError
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        This function checks the type of the provided `device` and returns the
+        appropriate context manager for that device.
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            contextmanager: A context manager specific to the device type.
+        """
+        raise NotImplementedError
+
+    def is_available(self) -> bool:
+        """
+        Check if the specific accelerator is available.
+        Returns:
+            bool: False by default, overridden by subclasses.
+        """
+        return False
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        raise NotImplementedError
+
+    def get_unique_id(self) -> str:
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: Done.
created_at: 2025-04-29T02:26:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,90 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorInfoHolder(Communicator):
+    """
+    The driver maintains a communicator instance but does not participate
+    in the group. However, it still needs to store certain information,
+    such as world size and actor details.
+
+    Since the driver may not have a GPU (or other accelerators), it could
+    mistakenly select the wrong `TorchDeviceManager`
+    (e.g., `CPUTorchDeviceManager`), causing `get_communicator`
+    to return `None`.
+
+    Because the driver does not actively join the group and only uses the
+    communicator for metadata storage, Introduced a hardware-independent
+    class specifically for the driver.
path: python/ray/experimental/channel/communicator_info_holder.py
user: login: hipudding
body: Done.
created_at: 2025-04-29T02:26:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
+    """
+
+    _torch_module_name = None
+    _instance = None
+    _communicator_cls = None
+    _torch_mod = None
+
+    @classmethod
+    def get(cls):
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+        if cls._instance is not None:
+            return cls._instance
+
+        with _accelerator_context_lock:
+            # Check _instance again to avoid creating a duplicate instance
+            # if it was initialized while waiting for the lock.
+            if cls._instance is None:
+                # Not registrey yet.
path: python/ray/experimental/channel/accelerator_context.py
user: login: noemotiovon
body: # Not registered yet.
created_at: 2025-04-29T06:56:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
+    """
+
+    _torch_module_name = None
+    _instance = None
+    _communicator_cls = None
+    _torch_mod = None
+
+    @classmethod
+    def get(cls):
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+        if cls._instance is not None:
+            return cls._instance
+
+        with _accelerator_context_lock:
+            # Check _instance again to avoid creating a duplicate instance
+            # if it was initialized while waiting for the lock.
+            if cls._instance is None:
+                # Not registrey yet.
path: python/ray/experimental/channel/accelerator_context.py
user: login: hipudding
body: Got it. Big enough.
created_at: 2025-04-29T07:46:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
path: python/ray/experimental/channel/accelerator_context.py
user: login: stephanie-wang
body: What is "it" in this comment?
created_at: 2025-05-07T00:50:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
+    """
+
+    _torch_module_name = None
+    _instance = None
+    _communicator_cls = None
+    _torch_mod = None
path: python/ray/experimental/channel/accelerator_context.py
user: login: stephanie-wang
body: Type hints and comments.
created_at: 2025-05-07T00:51:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
+    """
+
+    _torch_module_name = None
+    _instance = None
+    _communicator_cls = None
+    _torch_mod = None
+
+    @classmethod
+    def get(cls):
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+        if cls._instance is not None:
+            return cls._instance
+
+        with _accelerator_context_lock:
+            # Check _instance again to avoid creating a duplicate instance
+            # if it was initialized while waiting for the lock.
+            if cls._instance is None:
+                # Not registrey yet.
+                if cls._torch_mod is None:
+                    if len(ray.get_gpu_ids()) > 0:
+                        from ray.experimental.channel.nccl_group import _NcclGroup
+
+                        cls._set_context("cuda", _NcclGroup)
path: python/ray/experimental/channel/accelerator_context.py
user: login: stephanie-wang
body: Why use `_set_context` instead of just passing these args to `cls()` constructor?
created_at: 2025-05-07T00:53:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
+    """
+
+    _torch_module_name = None
+    _instance = None
+    _communicator_cls = None
+    _torch_mod = None
+
+    @classmethod
+    def get(cls):
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+        if cls._instance is not None:
+            return cls._instance
+
+        with _accelerator_context_lock:
+            # Check _instance again to avoid creating a duplicate instance
+            # if it was initialized while waiting for the lock.
+            if cls._instance is None:
+                # Not registrey yet.
+                if cls._torch_mod is None:
+                    if len(ray.get_gpu_ids()) > 0:
+                        from ray.experimental.channel.nccl_group import _NcclGroup
+
+                        cls._set_context("cuda", _NcclGroup)
+                    else:
+                        from ray.experimental.channel.cpu_communicator import (
+                            CPUCommunicator,
+                        )
+
+                        cls._set_context("cpu", CPUCommunicator)
+                cls._instance = cls()
+
+        return cls._instance
+
+    @classmethod
+    def _set_context(cls, name: str, communicator: Communicator):
+        """
+        Registers the accelerator's name and communicator class.
+
+        Args:
+            name: The name of the device module under torch (e.g., 'cuda', 'cpu').
+            communicator: The communicator class associated with the device.
+        """
+        cls._torch_module_name = name
+        cls._communicator_cls = communicator
+        # Save the torch module corresponding to the device module name.
+        cls._torch_mod = importlib.import_module(f"torch.{cls._torch_module_name}")
+
+    def get_default_device(self) -> "torch.device":
+        """
+        Returns the default device used by the compiled graph.
+
+        Currently, the default device for the compiled graph is the first visible
+        device. By default, it returns the device with logical ID 0.
+
+        Returns:
+            torch.device: The default device.
+        """
+        import torch
+
+        if self._torch_module_name == "cpu":
+            return torch.device("cpu")
+
+        return torch.device(f"{self._torch_module_name}:0")
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        There is no device context for CPU, returning a nullcontext.
+
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            device_context: A context manager specific to the device type.
+        """
+        if device.type == "cpu":
+            return nullcontext()
+
+        return self._torch_mod.device(device)
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        return self._torch_mod.current_stream()
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        return self._torch_mod.Event()
+
+    def generate_communicator_id(self) -> str:
+        """
+        Generates a communication identifier for communication group.
+        """
+        return self._communicator_cls.generate_communicator_id()
+
+    def get_communicator(self, *args, **kwargs):
path: python/ray/experimental/channel/accelerator_context.py
user: login: stephanie-wang
body: ```suggestion
    def create_communicator(self, *args, **kwargs):
```
created_at: 2025-05-07T00:56:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
+    """
+
+    _torch_module_name = None
+    _instance = None
+    _communicator_cls = None
+    _torch_mod = None
+
+    @classmethod
+    def get(cls):
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+        if cls._instance is not None:
+            return cls._instance
+
+        with _accelerator_context_lock:
+            # Check _instance again to avoid creating a duplicate instance
+            # if it was initialized while waiting for the lock.
+            if cls._instance is None:
+                # Not registrey yet.
+                if cls._torch_mod is None:
+                    if len(ray.get_gpu_ids()) > 0:
+                        from ray.experimental.channel.nccl_group import _NcclGroup
+
+                        cls._set_context("cuda", _NcclGroup)
+                    else:
+                        from ray.experimental.channel.cpu_communicator import (
+                            CPUCommunicator,
+                        )
+
+                        cls._set_context("cpu", CPUCommunicator)
+                cls._instance = cls()
+
+        return cls._instance
+
+    @classmethod
+    def _set_context(cls, name: str, communicator: Communicator):
+        """
+        Registers the accelerator's name and communicator class.
+
+        Args:
+            name: The name of the device module under torch (e.g., 'cuda', 'cpu').
+            communicator: The communicator class associated with the device.
+        """
+        cls._torch_module_name = name
+        cls._communicator_cls = communicator
+        # Save the torch module corresponding to the device module name.
+        cls._torch_mod = importlib.import_module(f"torch.{cls._torch_module_name}")
+
+    def get_default_device(self) -> "torch.device":
+        """
+        Returns the default device used by the compiled graph.
+
+        Currently, the default device for the compiled graph is the first visible
+        device. By default, it returns the device with logical ID 0.
+
+        Returns:
+            torch.device: The default device.
+        """
+        import torch
+
+        if self._torch_module_name == "cpu":
+            return torch.device("cpu")
+
+        return torch.device(f"{self._torch_module_name}:0")
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        There is no device context for CPU, returning a nullcontext.
+
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            device_context: A context manager specific to the device type.
+        """
+        if device.type == "cpu":
+            return nullcontext()
+
+        return self._torch_mod.device(device)
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        return self._torch_mod.current_stream()
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        return self._torch_mod.Event()
+
+    def generate_communicator_id(self) -> str:
+        """
+        Generates a communication identifier for communication group.
+        """
+        return self._communicator_cls.generate_communicator_id()
+
+    def get_communicator(self, *args, **kwargs):
+        """
+        Retrieves the communication group for collective operations.
path: python/ray/experimental/channel/accelerator_context.py
user: login: stephanie-wang
body: ```suggestion
        Creates a communication group for collective operations.
```
created_at: 2025-05-07T00:56:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -108,17 +107,17 @@ def recv(
 
     @property
     @abstractmethod
-    def recv_stream(self) -> Optional["cp.cuda.ExternalStream"]:
+    def recv_stream(self):
path: python/ray/experimental/channel/communicator.py
user: login: stephanie-wang
body: Can you add back the type hints?
created_at: 2025-05-07T00:56:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,106 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorHandle(Communicator):
+    """
+    A lightweight communicator handle used to store world size and actors
+    without providing actual communication capabilities.
+
+    This class is mainly used on the driver side where no real communicator
+    initialization or communication operations are needed. Most communication
+    methods are not implemented and will raise NotImplementedError if called.
path: python/ray/experimental/channel/communicator_handle.py
user: login: stephanie-wang
body: Why include the methods at all? Can we just create a separate class that doesn't inherit from Communicator and only includes the methods used by the driver?
created_at: 2025-05-07T00:58:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -675,14 +677,16 @@ def _do_destroy_communicator(self, group_id):
     # task loop running.
 
 
-def _do_check_has_gpu(self) -> bool:
-    return bool(ray.get_gpu_ids())
+def _do_check_has_communicator(self) -> bool:
+    return AcceleratorContext.get()._communicator_cls is not None
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: stephanie-wang
body: This method seems to check only if there is a communicator class but not if there is a communicator initialized? Should the function be renamed?
created_at: 2025-05-07T01:00:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -675,14 +677,16 @@ def _do_destroy_communicator(self, group_id):
     # task loop running.
 
 
-def _do_check_has_gpu(self) -> bool:
-    return bool(ray.get_gpu_ids())
+def _do_check_has_communicator(self) -> bool:
+    return AcceleratorContext.get()._communicator_cls is not None
 
 
-def _do_get_unique_nccl_id(self) -> bool:
-    from cupy.cuda import nccl
+def _do_set_custom_communicator(self, name: str, communicator: Communicator):
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: stephanie-wang
body: Rename to `do_register_accelerator_context`?
created_at: 2025-05-07T01:01:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -740,11 +744,24 @@ def _init_communicator(
         custom_communicator, CPUCommunicator
     )
 
-    has_gpus = ray.get(
-        [actor.__ray_call__.remote(_do_check_has_gpu) for actor in actors]
+    # set custom communicator on all actors
+    if AcceleratorContext._communicator_cls is not None:
+        ray.get(
+            [
+                actor.__ray_call__.remote(
+                    _do_set_custom_communicator,
+                    AcceleratorContext._torch_module_name,
+                    AcceleratorContext._communicator_cls,
+                )
+                for actor in actors
+            ]
+        )
+
+    has_acclerators = ray.get(
+        [actor.__ray_call__.remote(_do_check_has_communicator) for actor in actors]
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: stephanie-wang
body: A bit confusing that this is calling `_do_check_has_communicator` but the variable is called `has_accelerators`. Can't the actors have accelerators but not a communicator created yet?
created_at: 2025-05-07T01:02:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -94,68 +90,3 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
                 lambda self: ray.get_runtime_context().get_node_id()
             )
         )
-
-
-def get_cuda_devices() -> List["torch.device"]:
path: python/ray/experimental/channel/utils.py
user: login: stephanie-wang
body: Just checking - why was it that this code could be removed?
created_at: 2025-05-07T01:02:32Z
author_association: CONTRIBUTOR
diff_hunk: @@ -740,11 +744,24 @@ def _init_communicator(
         custom_communicator, CPUCommunicator
     )
 
-    has_gpus = ray.get(
-        [actor.__ray_call__.remote(_do_check_has_gpu) for actor in actors]
+    # set custom communicator on all actors
+    if AcceleratorContext._communicator_cls is not None:
+        ray.get(
+            [
+                actor.__ray_call__.remote(
+                    _do_set_custom_communicator,
+                    AcceleratorContext._torch_module_name,
+                    AcceleratorContext._communicator_cls,
+                )
+                for actor in actors
+            ]
+        )
+
+    has_acclerators = ray.get(
+        [actor.__ray_call__.remote(_do_check_has_communicator) for actor in actors]
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: stephanie-wang
body: Also I think this check is not exactly equivalent to the previous check - is it right that each actor could now have an AcceleratorContext but they could be for different devices? Do we have a unit test for this?
created_at: 2025-05-07T01:05:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -740,11 +744,24 @@ def _init_communicator(
         custom_communicator, CPUCommunicator
     )
 
-    has_gpus = ray.get(
-        [actor.__ray_call__.remote(_do_check_has_gpu) for actor in actors]
+    # set custom communicator on all actors
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: stephanie-wang
body: ```suggestion
    # Set custom communicator on all actors.
```
created_at: 2025-05-07T01:06:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -740,11 +744,24 @@ def _init_communicator(
         custom_communicator, CPUCommunicator
     )
 
-    has_gpus = ray.get(
-        [actor.__ray_call__.remote(_do_check_has_gpu) for actor in actors]
+    # set custom communicator on all actors
+    if AcceleratorContext._communicator_cls is not None:
+        ray.get(
+            [
+                actor.__ray_call__.remote(
+                    _do_set_custom_communicator,
+                    AcceleratorContext._torch_module_name,
+                    AcceleratorContext._communicator_cls,
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: stephanie-wang
body: Is there a way we can take in the `_communicator_cls` and `_torch_module_name` as arguments to this function? That makes it explicit what the dependencies of this function are, instead of passing the dependencies through the accelerator context.
created_at: 2025-05-07T01:08:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -740,11 +744,24 @@ def _init_communicator(
         custom_communicator, CPUCommunicator
     )
 
-    has_gpus = ray.get(
-        [actor.__ray_call__.remote(_do_check_has_gpu) for actor in actors]
+    # set custom communicator on all actors
+    if AcceleratorContext._communicator_cls is not None:
+        ray.get(
+            [
+                actor.__ray_call__.remote(
+                    _do_set_custom_communicator,
+                    AcceleratorContext._torch_module_name,
+                    AcceleratorContext._communicator_cls,
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: hipudding
body: > Is there a way we can take in the `_communicator_cls` and `_torch_module_name` as arguments to this function? That makes it explicit what the dependencies of this function are, instead of passing the dependencies through the accelerator context.

Are you suggesting that `_communicator_cls` and `_torch_module_name` should be passed as parameters to the `_init_communicator` function?
created_at: 2025-05-07T06:32:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -740,11 +744,24 @@ def _init_communicator(
         custom_communicator, CPUCommunicator
     )
 
-    has_gpus = ray.get(
-        [actor.__ray_call__.remote(_do_check_has_gpu) for actor in actors]
+    # set custom communicator on all actors
+    if AcceleratorContext._communicator_cls is not None:
+        ray.get(
+            [
+                actor.__ray_call__.remote(
+                    _do_set_custom_communicator,
+                    AcceleratorContext._torch_module_name,
+                    AcceleratorContext._communicator_cls,
+                )
+                for actor in actors
+            ]
+        )
+
+    has_acclerators = ray.get(
+        [actor.__ray_call__.remote(_do_check_has_communicator) for actor in actors]
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: hipudding
body: You're right. I will fix this.
created_at: 2025-05-07T07:17:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -94,68 +90,3 @@ def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str:
                 lambda self: ray.get_runtime_context().get_node_id()
             )
         )
-
-
-def get_cuda_devices() -> List["torch.device"]:
path: python/ray/experimental/channel/utils.py
user: login: hipudding
body: This function returns the logical IDs of the CUDA devices visible to the current worker. Currently, in CG, the device at index 0 is used as the default device, and the logical ID of the device at index 0 is also 0, so I believe calling get_cuda_devices is unnecessary. I think this part needs to be double-checked by @ruisearch42.
created_at: 2025-05-07T07:36:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -108,17 +107,17 @@ def recv(
 
     @property
     @abstractmethod
-    def recv_stream(self) -> Optional["cp.cuda.ExternalStream"]:
+    def recv_stream(self):
path: python/ray/experimental/channel/communicator.py
user: login: hipudding
body: Sorry, I don’t think we should specify a type hint here, because the stream context types differ across devices — for example, torch.cuda.StreamContext and torch.npu.StreamContext.
created_at: 2025-05-07T07:38:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,106 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorHandle(Communicator):
+    """
+    A lightweight communicator handle used to store world size and actors
+    without providing actual communication capabilities.
+
+    This class is mainly used on the driver side where no real communicator
+    initialization or communication operations are needed. Most communication
+    methods are not implemented and will raise NotImplementedError if called.
path: python/ray/experimental/channel/communicator_handle.py
user: login: hipudding
body: Yes, I agree that inheriting from Communicator may not be necessary. However, these Communicator handlers are stored in a Dict within the ChannelContext, where type hints are enforced. If we skip type checking here, wouldn’t that introduce a risk of type mismatches? For example, in future updates, someone might modify the Communicator base class but forget to update the corresponding Communicator handlers.
created_at: 2025-05-07T07:43:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
+    """
+
+    _torch_module_name = None
+    _instance = None
+    _communicator_cls = None
+    _torch_mod = None
+
+    @classmethod
+    def get(cls):
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+        if cls._instance is not None:
+            return cls._instance
+
+        with _accelerator_context_lock:
+            # Check _instance again to avoid creating a duplicate instance
+            # if it was initialized while waiting for the lock.
+            if cls._instance is None:
+                # Not registrey yet.
+                if cls._torch_mod is None:
+                    if len(ray.get_gpu_ids()) > 0:
+                        from ray.experimental.channel.nccl_group import _NcclGroup
+
+                        cls._set_context("cuda", _NcclGroup)
path: python/ray/experimental/channel/accelerator_context.py
user: login: hipudding
body: There are two steps here:
	1.	Set class attributes, including `_torch_module_name` and `_communicator_cls`. These attributes can be configured via `register_accelerator_context`; if not set, the system will automatically detect whether to use GPU or CPU.
	2.	Instantiate the AcceleratorContext object based on these class attributes.

If we were to use the constructor here, we would need to introduce global variables to store `_torch_module_name` and `_communicator_cls`.

Which approach do you think is better?
created_at: 2025-05-07T08:56:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
+    """
+
+    _torch_module_name = None
+    _instance = None
+    _communicator_cls = None
+    _torch_mod = None
+
+    @classmethod
+    def get(cls):
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+        if cls._instance is not None:
+            return cls._instance
+
+        with _accelerator_context_lock:
+            # Check _instance again to avoid creating a duplicate instance
+            # if it was initialized while waiting for the lock.
+            if cls._instance is None:
+                # Not registrey yet.
+                if cls._torch_mod is None:
+                    if len(ray.get_gpu_ids()) > 0:
+                        from ray.experimental.channel.nccl_group import _NcclGroup
+
+                        cls._set_context("cuda", _NcclGroup)
path: python/ray/experimental/channel/accelerator_context.py
user: login: stephanie-wang
body: Hmm the current code is brittle because if you call `AcceleratorContext._set_context`, it will modify a previously created singleton returned by `AcceleratorContext.get()` as well. It would be better if you can make it so that calling `AcceleratorContext.get()` returns an instance whose attributes are not tied to the class attributes. This make it a bit nicer for testing or if we eventually want to support multiple contexts in the same process.

I don't have a strong preference on how you should do that, although I agree it would be nice to not need global variables for the torch_module_name and communicator_cls. You could follow the [`DataContext` pattern](https://github.com/ray-project/ray/blob/3530f8e43d052e7ec983813d5f66d562aa3eba60/python/ray/data/context.py#L232), which stores the defaults on the class but uses `dataclass` and a global default context for the current instance.
created_at: 2025-05-07T17:54:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
+    """
+
+    _torch_module_name = None
+    _instance = None
+    _communicator_cls = None
+    _torch_mod = None
+
+    @classmethod
+    def get(cls):
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+        if cls._instance is not None:
+            return cls._instance
+
+        with _accelerator_context_lock:
+            # Check _instance again to avoid creating a duplicate instance
+            # if it was initialized while waiting for the lock.
+            if cls._instance is None:
+                # Not registrey yet.
path: python/ray/experimental/channel/accelerator_context.py
user: login: stephanie-wang
body: ```suggestion
                # Not registered yet.
```
created_at: 2025-05-07T17:55:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,106 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorHandle(Communicator):
+    """
+    A lightweight communicator handle used to store world size and actors
+    without providing actual communication capabilities.
+
+    This class is mainly used on the driver side where no real communicator
+    initialization or communication operations are needed. Most communication
+    methods are not implemented and will raise NotImplementedError if called.
path: python/ray/experimental/channel/communicator_handle.py
user: login: stephanie-wang
body: We can have separate fields in the context for CommunicatorHandle vs Communicator. This would actually be better because it's more explicit.
created_at: 2025-05-07T17:57:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -740,11 +744,24 @@ def _init_communicator(
         custom_communicator, CPUCommunicator
     )
 
-    has_gpus = ray.get(
-        [actor.__ray_call__.remote(_do_check_has_gpu) for actor in actors]
+    # set custom communicator on all actors
+    if AcceleratorContext._communicator_cls is not None:
+        ray.get(
+            [
+                actor.__ray_call__.remote(
+                    _do_set_custom_communicator,
+                    AcceleratorContext._torch_module_name,
+                    AcceleratorContext._communicator_cls,
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: stephanie-wang
body: Yes.
created_at: 2025-05-07T17:59:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,106 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple
+import ray
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from ray.experimental.util.types import ReduceOp
+
+if TYPE_CHECKING:
+    import torch
+
+
+class CommunicatorHandle(Communicator):
+    """
+    A lightweight communicator handle used to store world size and actors
+    without providing actual communication capabilities.
+
+    This class is mainly used on the driver side where no real communicator
+    initialization or communication operations are needed. Most communication
+    methods are not implemented and will raise NotImplementedError if called.
path: python/ray/experimental/channel/communicator_handle.py
user: login: hipudding
body: This still doesn’t seem to work, because when driver get the output_channel from worker https://github.com/ray-project/ray/blob/262af06532209c4dd81fe2046e29dab5af91bc9c/python/ray/dag/compiled_dag_node.py#L1647 , it uses the driver’s Communicator to get the ranks of the reader and writer in the channel. If we distinguish between CommunicatorHandle and Communicator, we would need to check whether we’re in the driver process, which I think adds complexity.

Can we put it simply, just skip fetching the group object when calling in driver, since even if it is fetched, it won’t be used.

created_at: 2025-05-08T12:25:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
path: python/ray/experimental/channel/accelerator_context.py
user: login: hipudding
body: Rewrite the class comment.
created_at: 2025-05-09T06:20:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
+    """
+
+    _torch_module_name = None
+    _instance = None
+    _communicator_cls = None
+    _torch_mod = None
path: python/ray/experimental/channel/accelerator_context.py
user: login: hipudding
body: Done.
created_at: 2025-05-09T06:20:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
+    """
+
+    _torch_module_name = None
+    _instance = None
+    _communicator_cls = None
+    _torch_mod = None
+
+    @classmethod
+    def get(cls):
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+        if cls._instance is not None:
+            return cls._instance
+
+        with _accelerator_context_lock:
+            # Check _instance again to avoid creating a duplicate instance
+            # if it was initialized while waiting for the lock.
+            if cls._instance is None:
+                # Not registrey yet.
+                if cls._torch_mod is None:
+                    if len(ray.get_gpu_ids()) > 0:
+                        from ray.experimental.channel.nccl_group import _NcclGroup
+
+                        cls._set_context("cuda", _NcclGroup)
path: python/ray/experimental/channel/accelerator_context.py
user: login: hipudding
body: Done.
created_at: 2025-05-09T06:20:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
+    """
+
+    _torch_module_name = None
+    _instance = None
+    _communicator_cls = None
+    _torch_mod = None
+
+    @classmethod
+    def get(cls):
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+        if cls._instance is not None:
+            return cls._instance
+
+        with _accelerator_context_lock:
+            # Check _instance again to avoid creating a duplicate instance
+            # if it was initialized while waiting for the lock.
+            if cls._instance is None:
+                # Not registrey yet.
+                if cls._torch_mod is None:
+                    if len(ray.get_gpu_ids()) > 0:
+                        from ray.experimental.channel.nccl_group import _NcclGroup
+
+                        cls._set_context("cuda", _NcclGroup)
+                    else:
+                        from ray.experimental.channel.cpu_communicator import (
+                            CPUCommunicator,
+                        )
+
+                        cls._set_context("cpu", CPUCommunicator)
+                cls._instance = cls()
+
+        return cls._instance
+
+    @classmethod
+    def _set_context(cls, name: str, communicator: Communicator):
+        """
+        Registers the accelerator's name and communicator class.
+
+        Args:
+            name: The name of the device module under torch (e.g., 'cuda', 'cpu').
+            communicator: The communicator class associated with the device.
+        """
+        cls._torch_module_name = name
+        cls._communicator_cls = communicator
+        # Save the torch module corresponding to the device module name.
+        cls._torch_mod = importlib.import_module(f"torch.{cls._torch_module_name}")
+
+    def get_default_device(self) -> "torch.device":
+        """
+        Returns the default device used by the compiled graph.
+
+        Currently, the default device for the compiled graph is the first visible
+        device. By default, it returns the device with logical ID 0.
+
+        Returns:
+            torch.device: The default device.
+        """
+        import torch
+
+        if self._torch_module_name == "cpu":
+            return torch.device("cpu")
+
+        return torch.device(f"{self._torch_module_name}:0")
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        There is no device context for CPU, returning a nullcontext.
+
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            device_context: A context manager specific to the device type.
+        """
+        if device.type == "cpu":
+            return nullcontext()
+
+        return self._torch_mod.device(device)
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        return self._torch_mod.current_stream()
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        return self._torch_mod.Event()
+
+    def generate_communicator_id(self) -> str:
+        """
+        Generates a communication identifier for communication group.
+        """
+        return self._communicator_cls.generate_communicator_id()
+
+    def get_communicator(self, *args, **kwargs):
path: python/ray/experimental/channel/accelerator_context.py
user: login: hipudding
body: Done.
created_at: 2025-05-09T06:21:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
+    """
+
+    _torch_module_name = None
+    _instance = None
+    _communicator_cls = None
+    _torch_mod = None
+
+    @classmethod
+    def get(cls):
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+        if cls._instance is not None:
+            return cls._instance
+
+        with _accelerator_context_lock:
+            # Check _instance again to avoid creating a duplicate instance
+            # if it was initialized while waiting for the lock.
+            if cls._instance is None:
+                # Not registrey yet.
+                if cls._torch_mod is None:
+                    if len(ray.get_gpu_ids()) > 0:
+                        from ray.experimental.channel.nccl_group import _NcclGroup
+
+                        cls._set_context("cuda", _NcclGroup)
+                    else:
+                        from ray.experimental.channel.cpu_communicator import (
+                            CPUCommunicator,
+                        )
+
+                        cls._set_context("cpu", CPUCommunicator)
+                cls._instance = cls()
+
+        return cls._instance
+
+    @classmethod
+    def _set_context(cls, name: str, communicator: Communicator):
+        """
+        Registers the accelerator's name and communicator class.
+
+        Args:
+            name: The name of the device module under torch (e.g., 'cuda', 'cpu').
+            communicator: The communicator class associated with the device.
+        """
+        cls._torch_module_name = name
+        cls._communicator_cls = communicator
+        # Save the torch module corresponding to the device module name.
+        cls._torch_mod = importlib.import_module(f"torch.{cls._torch_module_name}")
+
+    def get_default_device(self) -> "torch.device":
+        """
+        Returns the default device used by the compiled graph.
+
+        Currently, the default device for the compiled graph is the first visible
+        device. By default, it returns the device with logical ID 0.
+
+        Returns:
+            torch.device: The default device.
+        """
+        import torch
+
+        if self._torch_module_name == "cpu":
+            return torch.device("cpu")
+
+        return torch.device(f"{self._torch_module_name}:0")
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        There is no device context for CPU, returning a nullcontext.
+
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            device_context: A context manager specific to the device type.
+        """
+        if device.type == "cpu":
+            return nullcontext()
+
+        return self._torch_mod.device(device)
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        return self._torch_mod.current_stream()
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        return self._torch_mod.Event()
+
+    def generate_communicator_id(self) -> str:
+        """
+        Generates a communication identifier for communication group.
+        """
+        return self._communicator_cls.generate_communicator_id()
+
+    def get_communicator(self, *args, **kwargs):
+        """
+        Retrieves the communication group for collective operations.
path: python/ray/experimental/channel/accelerator_context.py
user: login: hipudding
body: Done.
created_at: 2025-05-09T06:21:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -675,14 +677,16 @@ def _do_destroy_communicator(self, group_id):
     # task loop running.
 
 
-def _do_check_has_gpu(self) -> bool:
-    return bool(ray.get_gpu_ids())
+def _do_check_has_communicator(self) -> bool:
+    return AcceleratorContext.get()._communicator_cls is not None
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: hipudding
body: Done.
created_at: 2025-05-09T06:21:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -675,14 +677,16 @@ def _do_destroy_communicator(self, group_id):
     # task loop running.
 
 
-def _do_check_has_gpu(self) -> bool:
-    return bool(ray.get_gpu_ids())
+def _do_check_has_communicator(self) -> bool:
+    return AcceleratorContext.get()._communicator_cls is not None
 
 
-def _do_get_unique_nccl_id(self) -> bool:
-    from cupy.cuda import nccl
+def _do_set_custom_communicator(self, name: str, communicator: Communicator):
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: hipudding
body: Done.
created_at: 2025-05-09T06:21:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -740,11 +744,24 @@ def _init_communicator(
         custom_communicator, CPUCommunicator
     )
 
-    has_gpus = ray.get(
-        [actor.__ray_call__.remote(_do_check_has_gpu) for actor in actors]
+    # set custom communicator on all actors
+    if AcceleratorContext._communicator_cls is not None:
+        ray.get(
+            [
+                actor.__ray_call__.remote(
+                    _do_set_custom_communicator,
+                    AcceleratorContext._torch_module_name,
+                    AcceleratorContext._communicator_cls,
+                )
+                for actor in actors
+            ]
+        )
+
+    has_acclerators = ray.get(
+        [actor.__ray_call__.remote(_do_check_has_communicator) for actor in actors]
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: hipudding
body: This has been changed to check whether all actors use the same accelerator, to ensure consistency across all actors. A test case added.
created_at: 2025-05-09T06:22:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -740,11 +744,24 @@ def _init_communicator(
         custom_communicator, CPUCommunicator
     )
 
-    has_gpus = ray.get(
-        [actor.__ray_call__.remote(_do_check_has_gpu) for actor in actors]
+    # set custom communicator on all actors
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: hipudding
body: Done.
created_at: 2025-05-09T06:22:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -740,11 +744,24 @@ def _init_communicator(
         custom_communicator, CPUCommunicator
     )
 
-    has_gpus = ray.get(
-        [actor.__ray_call__.remote(_do_check_has_gpu) for actor in actors]
+    # set custom communicator on all actors
+    if AcceleratorContext._communicator_cls is not None:
+        ray.get(
+            [
+                actor.__ray_call__.remote(
+                    _do_set_custom_communicator,
+                    AcceleratorContext._torch_module_name,
+                    AcceleratorContext._communicator_cls,
path: python/ray/experimental/channel/torch_tensor_nccl_channel.py
user: login: hipudding
body: Done.
created_at: 2025-05-09T06:23:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+_accelerator_context_lock = threading.Lock()
+
+
+class AcceleratorContext:
+    """
+    Provides a general interface for the Accelerator, including stream, event,
+    and device management.
+
+    This is a singleton class. Before using it, you must register the accelerator's
+    device module name and the Communicator class. If nothing is registered,
+    it will automatically fall back to GPU or CPU.
+    """
+
+    _torch_module_name = None
+    _instance = None
+    _communicator_cls = None
+    _torch_mod = None
+
+    @classmethod
+    def get(cls):
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+        if cls._instance is not None:
+            return cls._instance
+
+        with _accelerator_context_lock:
+            # Check _instance again to avoid creating a duplicate instance
+            # if it was initialized while waiting for the lock.
+            if cls._instance is None:
+                # Not registrey yet.
path: python/ray/experimental/channel/accelerator_context.py
user: login: hipudding
body: This part of code has been rewirte.
created_at: 2025-05-09T06:23:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,195 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This eincludes stram management, event creation, device context control,
path: python/ray/experimental/channel/accelerator_context.py
user: login: stephanie-wang
body: ```suggestion
    This includes stream management, event creation, device context control,
```
created_at: 2025-05-13T18:05:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,195 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This eincludes stram management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, commumcator_cls: Type[Communicator]):
path: python/ray/experimental/channel/accelerator_context.py
user: login: stephanie-wang
body: ```suggestion
    def __init__(self, torch_module_name: str, communicator_cls: Type[Communicator]):
```
created_at: 2025-05-13T18:05:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,195 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This eincludes stram management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, commumcator_cls: Type[Communicator]):
+        """
+        Initializes an accelerator context with the specified torch device module
+        and communicator class.
+
+        Args:
+            torch_module_name: Name of the torch device module (e.g., "cuda", "cpu").
+            commumcator_cls: Class used to handle communication.
path: python/ray/experimental/channel/accelerator_context.py
user: login: stephanie-wang
body: ```suggestion
            communicator_cls: Class used to handle communication.
```
created_at: 2025-05-13T18:05:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,195 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This eincludes stram management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, commumcator_cls: Type[Communicator]):
+        """
+        Initializes an accelerator context with the specified torch device module
+        and communicator class.
+
+        Args:
+            torch_module_name: Name of the torch device module (e.g., "cuda", "cpu").
+            commumcator_cls: Class used to handle communication.
+        """
+
+        # The name of the torch module (e.g., 'cuda', 'npu')
+        self._torch_module_name: str = torch_module_name
+        # The Communicator class used to manage communication
+        self._communicator_cls: Type[Communicator] = commumcator_cls
+
+        # Save the torch module corresponding to the device module name.
+        self._torch_mod = importlib.import_module(f"torch.{torch_module_name}")
+
+        # Whether the accelerator is registered.
+        self._is_registered: bool = False
+
+    @staticmethod
+    def get() -> "AcceleratorContext":
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+
+        global _default_accelerator_context
+
+        with _accelerator_context_lock:
+            if _default_accelerator_context is None:
+                if len(ray.get_gpu_ids()) > 0:
+                    from ray.experimental.channel.nccl_group import _NcclGroup
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cuda", _NcclGroup
+                    )
+                else:
+                    from ray.experimental.channel.cpu_communicator import (
+                        CPUCommunicator,
+                    )
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cpu", CPUCommunicator
+                    )
+
+            return _default_accelerator_context
+
+    @staticmethod
+    def set(accelerator_context: "AcceleratorContext") -> None:
+        """
+        Sets the accelerator context to the default context.
path: python/ray/experimental/channel/accelerator_context.py
user: login: stephanie-wang
body: ```suggestion
        Overwrites the default accelerator context.
```
created_at: 2025-05-13T18:07:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,195 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This eincludes stram management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, commumcator_cls: Type[Communicator]):
+        """
+        Initializes an accelerator context with the specified torch device module
+        and communicator class.
+
+        Args:
+            torch_module_name: Name of the torch device module (e.g., "cuda", "cpu").
+            commumcator_cls: Class used to handle communication.
+        """
+
+        # The name of the torch module (e.g., 'cuda', 'npu')
+        self._torch_module_name: str = torch_module_name
+        # The Communicator class used to manage communication
+        self._communicator_cls: Type[Communicator] = commumcator_cls
+
+        # Save the torch module corresponding to the device module name.
+        self._torch_mod = importlib.import_module(f"torch.{torch_module_name}")
+
+        # Whether the accelerator is registered.
+        self._is_registered: bool = False
+
+    @staticmethod
+    def get() -> "AcceleratorContext":
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+
+        global _default_accelerator_context
+
+        with _accelerator_context_lock:
+            if _default_accelerator_context is None:
+                if len(ray.get_gpu_ids()) > 0:
+                    from ray.experimental.channel.nccl_group import _NcclGroup
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cuda", _NcclGroup
+                    )
+                else:
+                    from ray.experimental.channel.cpu_communicator import (
+                        CPUCommunicator,
+                    )
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cpu", CPUCommunicator
+                    )
+
+            return _default_accelerator_context
+
+    @staticmethod
+    def set(accelerator_context: "AcceleratorContext") -> None:
+        """
+        Sets the accelerator context to the default context.
+
+        Args:
+            accelerator_context: The context to register.
+        """
+        global _default_accelerator_context
+
+        # Accelerator context is registered.
+        accelerator_context._is_registered = True
+        _default_accelerator_context = accelerator_context
+
+    def get_default_device(self) -> "torch.device":
+        """
+        Returns the default device used by the compiled graph.
+
+        Currently, the default device for the compiled graph is the first visible
+        device. By default, it returns the device with logical ID 0.
+
+        Returns:
+            torch.device: The default device.
+        """
+        import torch
+
+        if self._torch_module_name == "cpu":
+            return torch.device("cpu")
+
+        return torch.device(f"{self._torch_module_name}:0")
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        There is no device context for CPU, returning a nullcontext.
+
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            device_context: A context manager specific to the device type.
+        """
+        if device.type == "cpu":
+            return nullcontext()
+
+        return self._torch_mod.device(device)
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        return self._torch_mod.current_stream()
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        return self._torch_mod.Event()
+
+    def generate_communicator_id(self) -> str:
+        """
+        Generates a communication identifier for communication group.
+        """
+        return self._communicator_cls.generate_communicator_id()
+
+    def create_communicator(self, *args, **kwargs) -> Communicator:
+        """
+        Creates a communication group for collective operations.
+        """
+        return self._communicator_cls(*args, **kwargs)
+
+    def get_module_name(self) -> str:
path: python/ray/experimental/channel/accelerator_context.py
user: login: stephanie-wang
body: Could you change these getter methods to use `@property`?
created_at: 2025-05-13T18:08:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,195 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This eincludes stram management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, commumcator_cls: Type[Communicator]):
+        """
+        Initializes an accelerator context with the specified torch device module
+        and communicator class.
+
+        Args:
+            torch_module_name: Name of the torch device module (e.g., "cuda", "cpu").
+            commumcator_cls: Class used to handle communication.
+        """
+
+        # The name of the torch module (e.g., 'cuda', 'npu')
+        self._torch_module_name: str = torch_module_name
+        # The Communicator class used to manage communication
+        self._communicator_cls: Type[Communicator] = commumcator_cls
+
+        # Save the torch module corresponding to the device module name.
+        self._torch_mod = importlib.import_module(f"torch.{torch_module_name}")
+
+        # Whether the accelerator is registered.
+        self._is_registered: bool = False
+
+    @staticmethod
+    def get() -> "AcceleratorContext":
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+
+        global _default_accelerator_context
+
+        with _accelerator_context_lock:
+            if _default_accelerator_context is None:
+                if len(ray.get_gpu_ids()) > 0:
+                    from ray.experimental.channel.nccl_group import _NcclGroup
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cuda", _NcclGroup
+                    )
+                else:
+                    from ray.experimental.channel.cpu_communicator import (
+                        CPUCommunicator,
+                    )
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cpu", CPUCommunicator
+                    )
+
+            return _default_accelerator_context
+
+    @staticmethod
+    def set(accelerator_context: "AcceleratorContext") -> None:
+        """
+        Sets the accelerator context to the default context.
+
+        Args:
+            accelerator_context: The context to register.
+        """
+        global _default_accelerator_context
+
+        # Accelerator context is registered.
+        accelerator_context._is_registered = True
+        _default_accelerator_context = accelerator_context
+
+    def get_default_device(self) -> "torch.device":
+        """
+        Returns the default device used by the compiled graph.
+
+        Currently, the default device for the compiled graph is the first visible
+        device. By default, it returns the device with logical ID 0.
+
+        Returns:
+            torch.device: The default device.
+        """
+        import torch
+
+        if self._torch_module_name == "cpu":
+            return torch.device("cpu")
+
+        return torch.device(f"{self._torch_module_name}:0")
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        There is no device context for CPU, returning a nullcontext.
+
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            device_context: A context manager specific to the device type.
+        """
+        if device.type == "cpu":
+            return nullcontext()
+
+        return self._torch_mod.device(device)
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        return self._torch_mod.current_stream()
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        return self._torch_mod.Event()
+
+    def generate_communicator_id(self) -> str:
+        """
+        Generates a communication identifier for communication group.
+        """
+        return self._communicator_cls.generate_communicator_id()
+
+    def create_communicator(self, *args, **kwargs) -> Communicator:
+        """
+        Creates a communication group for collective operations.
+        """
+        return self._communicator_cls(*args, **kwargs)
+
+    def get_module_name(self) -> str:
+        """
+        Gets the name of the torch module backing the accelerator.
+        """
+        return self._torch_module_name
+
+    def get_registed_module_name(self) -> Optional[str]:
+        """
+        Returns the name of the module if registered. Otherwise, returns None.
+        """
+        if self._is_registered:
+            return self._torch_module_name
+        return None
+
+    def get_registed_communicator_cls(self) -> Optional[Type[Communicator]]:
+        """
+        Returns the communicator class if registered. Otherwise, returns None.
+        """
+        if self._is_registered:
+            return self._communicator_cls
+        return None
+
+    def get_accelerator_count(self) -> int:
+        """
+        Returns the number of accelerators assigned by ray.
+        """
+        if self._torch_module_name == "cuda":
+            return len(ray.get_gpu_ids())
+        else:
+            accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
+            return len(accelerator_ids.get(self._torch_module_name.upper(), []))
+
+
+def register_accelerator_context(
+    torch_module_name: str, communicator_cls: Type[Communicator]
+):
+    """
+    Registers the accelerator context with the specified device type and communicator.
+
+    Args:
+        torch_module_name: The name of the device module under torch.
+        communicator: The communicator class associated with the device.
+    """
+    accleerator_context = AcceleratorContext(torch_module_name, communicator_cls)
path: python/ray/experimental/channel/accelerator_context.py
user: login: stephanie-wang
body: ```suggestion
    accelerator_context = AcceleratorContext(torch_module_name, communicator_cls)
```
created_at: 2025-05-13T18:18:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,195 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This eincludes stram management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, commumcator_cls: Type[Communicator]):
+        """
+        Initializes an accelerator context with the specified torch device module
+        and communicator class.
+
+        Args:
+            torch_module_name: Name of the torch device module (e.g., "cuda", "cpu").
+            commumcator_cls: Class used to handle communication.
+        """
+
+        # The name of the torch module (e.g., 'cuda', 'npu')
+        self._torch_module_name: str = torch_module_name
+        # The Communicator class used to manage communication
+        self._communicator_cls: Type[Communicator] = commumcator_cls
+
+        # Save the torch module corresponding to the device module name.
+        self._torch_mod = importlib.import_module(f"torch.{torch_module_name}")
+
+        # Whether the accelerator is registered.
+        self._is_registered: bool = False
+
+    @staticmethod
+    def get() -> "AcceleratorContext":
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+
+        global _default_accelerator_context
+
+        with _accelerator_context_lock:
+            if _default_accelerator_context is None:
+                if len(ray.get_gpu_ids()) > 0:
+                    from ray.experimental.channel.nccl_group import _NcclGroup
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cuda", _NcclGroup
+                    )
+                else:
+                    from ray.experimental.channel.cpu_communicator import (
+                        CPUCommunicator,
+                    )
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cpu", CPUCommunicator
+                    )
+
+            return _default_accelerator_context
+
+    @staticmethod
+    def set(accelerator_context: "AcceleratorContext") -> None:
+        """
+        Sets the accelerator context to the default context.
+
+        Args:
+            accelerator_context: The context to register.
+        """
+        global _default_accelerator_context
+
+        # Accelerator context is registered.
+        accelerator_context._is_registered = True
+        _default_accelerator_context = accelerator_context
+
+    def get_default_device(self) -> "torch.device":
+        """
+        Returns the default device used by the compiled graph.
+
+        Currently, the default device for the compiled graph is the first visible
+        device. By default, it returns the device with logical ID 0.
+
+        Returns:
+            torch.device: The default device.
+        """
+        import torch
+
+        if self._torch_module_name == "cpu":
+            return torch.device("cpu")
+
+        return torch.device(f"{self._torch_module_name}:0")
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        There is no device context for CPU, returning a nullcontext.
+
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            device_context: A context manager specific to the device type.
+        """
+        if device.type == "cpu":
+            return nullcontext()
+
+        return self._torch_mod.device(device)
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        return self._torch_mod.current_stream()
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        return self._torch_mod.Event()
+
+    def generate_communicator_id(self) -> str:
+        """
+        Generates a communication identifier for communication group.
+        """
+        return self._communicator_cls.generate_communicator_id()
+
+    def create_communicator(self, *args, **kwargs) -> Communicator:
+        """
+        Creates a communication group for collective operations.
+        """
+        return self._communicator_cls(*args, **kwargs)
+
+    def get_module_name(self) -> str:
+        """
+        Gets the name of the torch module backing the accelerator.
+        """
+        return self._torch_module_name
+
+    def get_registed_module_name(self) -> Optional[str]:
+        """
+        Returns the name of the module if registered. Otherwise, returns None.
+        """
+        if self._is_registered:
+            return self._torch_module_name
+        return None
path: python/ray/experimental/channel/accelerator_context.py
user: login: stephanie-wang
body: I found this a bit confusing: I think it is odd that you can call register multiple times and end up with multiple AcceleratorContexts that have `_is_registered=True`.

Instead, how about removing the `_is_registered` flag and just having `get_module_name`. You can create a `_global_custom_context` along with the `_default_accelerator_context`, and `AcceleratorContext.get()` will return the custom context if set else it returns the default context. This makes it more explicit which context is registered or not and whether the context is custom or not.
created_at: 2025-05-13T18:21:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,26 @@
+from typing import List
+import ray
+
+
+class CommunicatorHandle:
+    """
+    A lightweight communicator handle used to store actors in the communicator.
path: python/ray/experimental/channel/communicator_handle.py
user: login: stephanie-wang
body: ```suggestion
    A lightweight communicator handle used by the driver to store handles to the actors in the communicator.
```
created_at: 2025-05-13T18:22:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,195 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This eincludes stram management, event creation, device context control,
path: python/ray/experimental/channel/accelerator_context.py
user: login: hipudding
body: Done.
created_at: 2025-05-19T02:38:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,195 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This eincludes stram management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, commumcator_cls: Type[Communicator]):
path: python/ray/experimental/channel/accelerator_context.py
user: login: hipudding
body: Done.
created_at: 2025-05-19T02:38:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,195 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This eincludes stram management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, commumcator_cls: Type[Communicator]):
+        """
+        Initializes an accelerator context with the specified torch device module
+        and communicator class.
+
+        Args:
+            torch_module_name: Name of the torch device module (e.g., "cuda", "cpu").
+            commumcator_cls: Class used to handle communication.
path: python/ray/experimental/channel/accelerator_context.py
user: login: hipudding
body: Done.
created_at: 2025-05-19T02:38:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,195 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This eincludes stram management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, commumcator_cls: Type[Communicator]):
+        """
+        Initializes an accelerator context with the specified torch device module
+        and communicator class.
+
+        Args:
+            torch_module_name: Name of the torch device module (e.g., "cuda", "cpu").
+            commumcator_cls: Class used to handle communication.
+        """
+
+        # The name of the torch module (e.g., 'cuda', 'npu')
+        self._torch_module_name: str = torch_module_name
+        # The Communicator class used to manage communication
+        self._communicator_cls: Type[Communicator] = commumcator_cls
+
+        # Save the torch module corresponding to the device module name.
+        self._torch_mod = importlib.import_module(f"torch.{torch_module_name}")
+
+        # Whether the accelerator is registered.
+        self._is_registered: bool = False
+
+    @staticmethod
+    def get() -> "AcceleratorContext":
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+
+        global _default_accelerator_context
+
+        with _accelerator_context_lock:
+            if _default_accelerator_context is None:
+                if len(ray.get_gpu_ids()) > 0:
+                    from ray.experimental.channel.nccl_group import _NcclGroup
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cuda", _NcclGroup
+                    )
+                else:
+                    from ray.experimental.channel.cpu_communicator import (
+                        CPUCommunicator,
+                    )
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cpu", CPUCommunicator
+                    )
+
+            return _default_accelerator_context
+
+    @staticmethod
+    def set(accelerator_context: "AcceleratorContext") -> None:
+        """
+        Sets the accelerator context to the default context.
path: python/ray/experimental/channel/accelerator_context.py
user: login: hipudding
body: Done.
created_at: 2025-05-19T02:39:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,195 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This eincludes stram management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, commumcator_cls: Type[Communicator]):
+        """
+        Initializes an accelerator context with the specified torch device module
+        and communicator class.
+
+        Args:
+            torch_module_name: Name of the torch device module (e.g., "cuda", "cpu").
+            commumcator_cls: Class used to handle communication.
+        """
+
+        # The name of the torch module (e.g., 'cuda', 'npu')
+        self._torch_module_name: str = torch_module_name
+        # The Communicator class used to manage communication
+        self._communicator_cls: Type[Communicator] = commumcator_cls
+
+        # Save the torch module corresponding to the device module name.
+        self._torch_mod = importlib.import_module(f"torch.{torch_module_name}")
+
+        # Whether the accelerator is registered.
+        self._is_registered: bool = False
+
+    @staticmethod
+    def get() -> "AcceleratorContext":
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+
+        global _default_accelerator_context
+
+        with _accelerator_context_lock:
+            if _default_accelerator_context is None:
+                if len(ray.get_gpu_ids()) > 0:
+                    from ray.experimental.channel.nccl_group import _NcclGroup
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cuda", _NcclGroup
+                    )
+                else:
+                    from ray.experimental.channel.cpu_communicator import (
+                        CPUCommunicator,
+                    )
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cpu", CPUCommunicator
+                    )
+
+            return _default_accelerator_context
+
+    @staticmethod
+    def set(accelerator_context: "AcceleratorContext") -> None:
+        """
+        Sets the accelerator context to the default context.
+
+        Args:
+            accelerator_context: The context to register.
+        """
+        global _default_accelerator_context
+
+        # Accelerator context is registered.
+        accelerator_context._is_registered = True
+        _default_accelerator_context = accelerator_context
+
+    def get_default_device(self) -> "torch.device":
+        """
+        Returns the default device used by the compiled graph.
+
+        Currently, the default device for the compiled graph is the first visible
+        device. By default, it returns the device with logical ID 0.
+
+        Returns:
+            torch.device: The default device.
+        """
+        import torch
+
+        if self._torch_module_name == "cpu":
+            return torch.device("cpu")
+
+        return torch.device(f"{self._torch_module_name}:0")
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        There is no device context for CPU, returning a nullcontext.
+
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            device_context: A context manager specific to the device type.
+        """
+        if device.type == "cpu":
+            return nullcontext()
+
+        return self._torch_mod.device(device)
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        return self._torch_mod.current_stream()
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        return self._torch_mod.Event()
+
+    def generate_communicator_id(self) -> str:
+        """
+        Generates a communication identifier for communication group.
+        """
+        return self._communicator_cls.generate_communicator_id()
+
+    def create_communicator(self, *args, **kwargs) -> Communicator:
+        """
+        Creates a communication group for collective operations.
+        """
+        return self._communicator_cls(*args, **kwargs)
+
+    def get_module_name(self) -> str:
path: python/ray/experimental/channel/accelerator_context.py
user: login: hipudding
body: Done.
created_at: 2025-05-19T02:39:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,195 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This eincludes stram management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, commumcator_cls: Type[Communicator]):
+        """
+        Initializes an accelerator context with the specified torch device module
+        and communicator class.
+
+        Args:
+            torch_module_name: Name of the torch device module (e.g., "cuda", "cpu").
+            commumcator_cls: Class used to handle communication.
+        """
+
+        # The name of the torch module (e.g., 'cuda', 'npu')
+        self._torch_module_name: str = torch_module_name
+        # The Communicator class used to manage communication
+        self._communicator_cls: Type[Communicator] = commumcator_cls
+
+        # Save the torch module corresponding to the device module name.
+        self._torch_mod = importlib.import_module(f"torch.{torch_module_name}")
+
+        # Whether the accelerator is registered.
+        self._is_registered: bool = False
+
+    @staticmethod
+    def get() -> "AcceleratorContext":
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+
+        global _default_accelerator_context
+
+        with _accelerator_context_lock:
+            if _default_accelerator_context is None:
+                if len(ray.get_gpu_ids()) > 0:
+                    from ray.experimental.channel.nccl_group import _NcclGroup
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cuda", _NcclGroup
+                    )
+                else:
+                    from ray.experimental.channel.cpu_communicator import (
+                        CPUCommunicator,
+                    )
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cpu", CPUCommunicator
+                    )
+
+            return _default_accelerator_context
+
+    @staticmethod
+    def set(accelerator_context: "AcceleratorContext") -> None:
+        """
+        Sets the accelerator context to the default context.
+
+        Args:
+            accelerator_context: The context to register.
+        """
+        global _default_accelerator_context
+
+        # Accelerator context is registered.
+        accelerator_context._is_registered = True
+        _default_accelerator_context = accelerator_context
+
+    def get_default_device(self) -> "torch.device":
+        """
+        Returns the default device used by the compiled graph.
+
+        Currently, the default device for the compiled graph is the first visible
+        device. By default, it returns the device with logical ID 0.
+
+        Returns:
+            torch.device: The default device.
+        """
+        import torch
+
+        if self._torch_module_name == "cpu":
+            return torch.device("cpu")
+
+        return torch.device(f"{self._torch_module_name}:0")
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        There is no device context for CPU, returning a nullcontext.
+
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            device_context: A context manager specific to the device type.
+        """
+        if device.type == "cpu":
+            return nullcontext()
+
+        return self._torch_mod.device(device)
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        return self._torch_mod.current_stream()
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        return self._torch_mod.Event()
+
+    def generate_communicator_id(self) -> str:
+        """
+        Generates a communication identifier for communication group.
+        """
+        return self._communicator_cls.generate_communicator_id()
+
+    def create_communicator(self, *args, **kwargs) -> Communicator:
+        """
+        Creates a communication group for collective operations.
+        """
+        return self._communicator_cls(*args, **kwargs)
+
+    def get_module_name(self) -> str:
+        """
+        Gets the name of the torch module backing the accelerator.
+        """
+        return self._torch_module_name
+
+    def get_registed_module_name(self) -> Optional[str]:
+        """
+        Returns the name of the module if registered. Otherwise, returns None.
+        """
+        if self._is_registered:
+            return self._torch_module_name
+        return None
+
+    def get_registed_communicator_cls(self) -> Optional[Type[Communicator]]:
+        """
+        Returns the communicator class if registered. Otherwise, returns None.
+        """
+        if self._is_registered:
+            return self._communicator_cls
+        return None
+
+    def get_accelerator_count(self) -> int:
+        """
+        Returns the number of accelerators assigned by ray.
+        """
+        if self._torch_module_name == "cuda":
+            return len(ray.get_gpu_ids())
+        else:
+            accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
+            return len(accelerator_ids.get(self._torch_module_name.upper(), []))
+
+
+def register_accelerator_context(
+    torch_module_name: str, communicator_cls: Type[Communicator]
+):
+    """
+    Registers the accelerator context with the specified device type and communicator.
+
+    Args:
+        torch_module_name: The name of the device module under torch.
+        communicator: The communicator class associated with the device.
+    """
+    accleerator_context = AcceleratorContext(torch_module_name, communicator_cls)
path: python/ray/experimental/channel/accelerator_context.py
user: login: hipudding
body: Done.
created_at: 2025-05-19T02:39:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,26 @@
+from typing import List
+import ray
+
+
+class CommunicatorHandle:
+    """
+    A lightweight communicator handle used to store actors in the communicator.
path: python/ray/experimental/channel/communicator_handle.py
user: login: hipudding
body: Done.
created_at: 2025-05-19T02:39:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,195 @@
+import threading
+import importlib
+import ray
+from typing import TYPE_CHECKING, Optional, Type
+from contextlib import nullcontext
+from ray.experimental.channel.communicator import Communicator
+
+if TYPE_CHECKING:
+    import torch
+
+# The accelerator context singleton on this process.
+_accelerator_context_lock = threading.Lock()
+_default_accelerator_context: Optional["AcceleratorContext"] = None
+
+
+class AcceleratorContext:
+    """
+    Provides a unified interface for managing different accelerator backends
+    This eincludes stram management, event creation, device context control,
+    and communicator support for distributed communication.
+    """
+
+    def __init__(self, torch_module_name: str, commumcator_cls: Type[Communicator]):
+        """
+        Initializes an accelerator context with the specified torch device module
+        and communicator class.
+
+        Args:
+            torch_module_name: Name of the torch device module (e.g., "cuda", "cpu").
+            commumcator_cls: Class used to handle communication.
+        """
+
+        # The name of the torch module (e.g., 'cuda', 'npu')
+        self._torch_module_name: str = torch_module_name
+        # The Communicator class used to manage communication
+        self._communicator_cls: Type[Communicator] = commumcator_cls
+
+        # Save the torch module corresponding to the device module name.
+        self._torch_mod = importlib.import_module(f"torch.{torch_module_name}")
+
+        # Whether the accelerator is registered.
+        self._is_registered: bool = False
+
+    @staticmethod
+    def get() -> "AcceleratorContext":
+        """
+        Returns the singleton instance of the accelerator context.
+
+        If a custom accelerator has been registered, initializes the context
+        based on the registration. Otherwise, selects an appropriate runtime
+        based on the available device (CUDA or CPU) and registers the
+        corresponding default communicator.
+
+        Returns:
+            AcceleratorContext: A singleton instance of the appropriate
+            runtime context.
+        """
+
+        global _default_accelerator_context
+
+        with _accelerator_context_lock:
+            if _default_accelerator_context is None:
+                if len(ray.get_gpu_ids()) > 0:
+                    from ray.experimental.channel.nccl_group import _NcclGroup
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cuda", _NcclGroup
+                    )
+                else:
+                    from ray.experimental.channel.cpu_communicator import (
+                        CPUCommunicator,
+                    )
+
+                    _default_accelerator_context = AcceleratorContext(
+                        "cpu", CPUCommunicator
+                    )
+
+            return _default_accelerator_context
+
+    @staticmethod
+    def set(accelerator_context: "AcceleratorContext") -> None:
+        """
+        Sets the accelerator context to the default context.
+
+        Args:
+            accelerator_context: The context to register.
+        """
+        global _default_accelerator_context
+
+        # Accelerator context is registered.
+        accelerator_context._is_registered = True
+        _default_accelerator_context = accelerator_context
+
+    def get_default_device(self) -> "torch.device":
+        """
+        Returns the default device used by the compiled graph.
+
+        Currently, the default device for the compiled graph is the first visible
+        device. By default, it returns the device with logical ID 0.
+
+        Returns:
+            torch.device: The default device.
+        """
+        import torch
+
+        if self._torch_module_name == "cpu":
+            return torch.device("cpu")
+
+        return torch.device(f"{self._torch_module_name}:0")
+
+    def get_device_context(self, device):
+        """
+        Retrieves the context manager for the specified accelerator device.
+        There is no device context for CPU, returning a nullcontext.
+
+        Args:
+            device (torch.device): The target device for which the context manager
+            is required.
+        Returns:
+            device_context: A context manager specific to the device type.
+        """
+        if device.type == "cpu":
+            return nullcontext()
+
+        return self._torch_mod.device(device)
+
+    def current_stream(self):
+        """
+        Retrieves the current execution stream for the accelerator device.
+        """
+        return self._torch_mod.current_stream()
+
+    def create_event(self):
+        """
+        Creates an event object for the accelerator device.
+        """
+        return self._torch_mod.Event()
+
+    def generate_communicator_id(self) -> str:
+        """
+        Generates a communication identifier for communication group.
+        """
+        return self._communicator_cls.generate_communicator_id()
+
+    def create_communicator(self, *args, **kwargs) -> Communicator:
+        """
+        Creates a communication group for collective operations.
+        """
+        return self._communicator_cls(*args, **kwargs)
+
+    def get_module_name(self) -> str:
+        """
+        Gets the name of the torch module backing the accelerator.
+        """
+        return self._torch_module_name
+
+    def get_registed_module_name(self) -> Optional[str]:
+        """
+        Returns the name of the module if registered. Otherwise, returns None.
+        """
+        if self._is_registered:
+            return self._torch_module_name
+        return None
path: python/ray/experimental/channel/accelerator_context.py
user: login: noemotiovon
body: I fully agree with your suggestion. I have implemented _global_custom_context following your design and provided a method is_accelerator_context_registered to check whether a context is registered. The code now looks much clearer. Thank you again for your valuable feedback!
created_at: 2025-05-19T03:12:56Z
author_association: CONTRIBUTOR

PR comments:
user: login: hipudding
created_at: 2025-03-04T06:16:16Z
body: @ruisearch42 Good Day. Could you please review this PR? Thanks.
author_association: CONTRIBUTOR
user: login: ruisearch42
created_at: 2025-03-05T04:35:08Z
body: Hi @hipudding , is this ready for review? btw, aDAG was renamed to Compiled Graph: https://docs.ray.io/en/latest/ray-core/compiled-graph/ray-compiled-graph.html
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-03-05T04:38:54Z
body: > Hi @hipudding , is this ready for review? btw, aDAG was renamed to Compiled Graph: https://docs.ray.io/en/latest/ray-core/compiled-graph/ray-compiled-graph.html

Yes, I think the main functionality of this PR is ready, I may make some minor adjustments later. Thanks for your review.
I'vd noticed that there's some test cases failed. I will fix them ASAP.
author_association: CONTRIBUTOR
user: login: ruisearch42
created_at: 2025-03-05T05:54:48Z
body: > Yes, I think the main functionality of this PR is ready, I may make some minor adjustments later. Thanks for your review. I'vd noticed that there's some test cases failed. I will fix them ASAP.

Thanks. I will review tomorrow!

author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-03-07T06:26:43Z
body: Hi @ruisearch42, Could you please review this PR again？

BTW, I fixed test case fail in doc_test_cgraph_nccl. It blocked this PR.
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-03-10T09:21:46Z
body: Bug when set transport='CPU'
serialization_context is not suite for NPU now. #50778
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-03-12T06:15:58Z
body: @ruisearch42  Morning.  Is this PR still under reviewing? Please let me know if there are any adjustments needed.
author_association: CONTRIBUTOR
user: login: ruisearch42
created_at: 2025-03-12T16:38:51Z
body: Hi @hipudding , see the discussions in https://github.com/ray-project/ray/pull/51305
We are breaking the dependency between Ray Core and other libraries, and the plan is to not depend on AIR.
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-03-13T01:03:21Z
body: > Hi @hipudding , see the discussions in #51305 We are breaking the dependency between Ray Core and other libraries, and the plan is to not depend on AIR.

Thanks. I will make some adjustments. 
I found this PR only provide a function to get default device. Can I add more functions for operating stream and event for different devices? ( Option1 in that discussions, which will make some duplicate code)
author_association: CONTRIBUTOR
user: login: noemotiovon
created_at: 2025-03-13T01:42:28Z
body: # vLLM on Ray benchmark(Ascend NPU)

## Test Environment

**python: 3.10**

**cann: 8.1.0**

**vllm: [vllm-project/vllm#14668](https://github.com/vllm-project/vllm/pull/14668)**

**ascend: main branch(0.7.3)**

**ray: [ray-project/ray#51032](https://github.com/ray-project/ray/pull/51032)**

**torch-npu: [2.5.1 rc1](https://github.com/vllm-project/vllm-ascend/blob/main/pta_install.sh)**

**pytorch: 2.5.1**

## Test Plan

The performance tests were conducted using vLLM's official Serving-Test script, [benchmark_serving.py](https://github.com/vllm-project/vllm/blob/main/benchmarks/benchmark_serving.py), under different scenarios, and the test results were recorded.

**Models: Qwen-2.5.0-0.5B/Qwen-2.5.0-14B**

**Scripts:**

```bash
# Serve
vllm serve Qwen/Qwen2.5-0.5B-Instruct \
  --pipeline-parallel-size 2 \
  --distributed-executor-backend ray \
  --trust-remote-code \
  --disable-log-stats \
  --disable-log-requests

# CLIENT
python3 benchmarks/benchmark_serving.py --backend vllm \
	--model Qwen/Qwen2.5-0.5B-Instruct \
	--dataset-name sharegpt \
	--dataset-path benchmarks/ShareGPT_V3_unfiltered_cleaned_split.json \
	--num-prompts 100 \
	--save-result \
	--result-dir benchmarks/result/ \
	--request-rate 4
```

This involves adjusting the parameters: `pipeline-parallel-size`, `model`, and `num-prompts`.

**vLLM Environment Variables:**

```bash
export VLLM_USE_RAY_SPMD_WORKER='1'
export VLLM_USE_RAY_COMPILED_DAG='1'
export VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM='1'
```



## Test Results

### **Qwen-2.5-0.5B**

| Test Scenarios                       | Throughput (tok/s) | Median TTFT (ms) | Median ITL (ms) |
| ------------------------------------ | ------------------ | ---------------- | --------------- |
| TP2 + Ray + CG disabled              | 867.49             | 110.95           | 43.81           |
| TP2 + Ray + CG                       | 858.43             | 67.93            | 45.46           |
| PP2 + MP                             | 417.00             | 72.91            | 50.36           |
| PP2 + Ray + CG disabled              | 780.89             | 79.29            | 46.91           |
| PP2 + Ray + CG + HCCL disabled       | 974.42             | 63.46            | 35.85           |
| PP2 + Ray + CG + HCCL                | 1006.72            | 60.40            | 34.81           |
| PP2 + Ray + CG + HCCL + OVERLAP COMM | 966.54             | 61.96            | 36.58           |

**Test Summary:**
For TP scenarios, there is no difference between enabling and disabling the CG feature.
**For PP scenarios, enabling CG + HCCL improves throughput by 29.1%, reduces first token latency by 23.8%, and decreases inter-token latency by 25.8% compared to disabling CG.**

### **Qwen-2.5-14B**

#### PP=2，num-prompts=100

| Test Scenarios                       | Throughput (tok/s) | Median TTFT (ms) | Median ITL (ms) |
| ------------------------------------ | ------------------ | ---------------- | --------------- |
| PP2 + MP                             | 547.93             | 129.24           | 75.23           |
| PP2 + Ray + CG disabled              | 581.58             | 140.40           | 78.35           |
| PP2 + Ray + CG + HCCL disabled       | 680.36             | 121.03           | 63.91           |
| PP2 + Ray + CG + HCCL                | 690.78             | 118.50           | 62.34           |
| PP2 + Ray + CG + HCCL + OVERLAP COMM | 664.87             | 126.62           | 64.84           |

**Test Summary:**
**For PP scenarios, enabling CG + HCCL improves throughput by 18.8%, reduces first token latency by 15.6%, and decreases inter-token latency by 20.4% compared to disabling CG.**

#### PP=4，num-prompts=500

| Test Scenarios                       | Throughput (tok/s) | Median TTFT (ms) | Median ITL (ms) |
| ------------------------------------ | ------------------ | ---------------- | --------------- |
| PP4 + Ray + CG disabled              | /                  | /                | /               |
| PP4 + Ray + CG + HCCL disabled       | 1237.83            | 134.52           | 71.74           |
| PP4 + Ray + CG + HCCL                | 1220.42            | 134.16           | 75.75           |
| PP4 + Ray + CG + HCCL + OVERLAP COMM | 1208.81            | 131.44           | 77.87           |
author_association: CONTRIBUTOR
user: login: ruisearch42
created_at: 2025-03-13T01:51:29Z
body: > I found this PR only provide a function to get default device. Can I add more functions for operating stream and event for different devices? ( Option1 in that discussions, which will make some duplicate code)

@hipudding I think we can start small and support GPU and NPU only, but having a general enough interface so that it is easily extensible.
author_association: CONTRIBUTOR
user: login: ruisearch42
created_at: 2025-03-13T01:53:25Z
body: @noemotiovon Thanks for sharing the benchmarks. The improvement looks pretty good! QQ, why enabling HCCL isn't better for PP=4?
author_association: CONTRIBUTOR
user: login: noemotiovon
created_at: 2025-03-13T01:59:21Z
body: @ruisearch42 ,
I also conducted tests on the GPU, and the results are consistent with those on the NPU:

| Test Scenarios                       | Throughput (tok/s) | Median TTFT (ms) | Median ITL (ms) |
| ------------------------------------ | ------------------ | ---------------- | --------------- |
| PP2 + Ray + CG disabled              | 1214.59            | 55.77            | 18.35           |
| PP2 + Ray + CG + NCCL disabled       | 1550.95            | 43.48            | 10.52           |
| PP2 + Ray + CG + NCCL                | 1546.07            | 43.90            | 11.15           |
| PP2 + Ray + CG + NCCL + OVERLAP COMM | 1546.66            | 42.23            | 11.18           |

**Test Summary:**
**For PP scenarios, enabling CG + NCCL improves throughput by 27.2%, reduces first token latency by 21.3%, and decreases inter-token latency by 39.2% compared to disabling CG.**
author_association: CONTRIBUTOR
user: login: noemotiovon
created_at: 2025-03-13T02:04:31Z
body: @ruisearch42 ,
Enabling NCCL and disabling NCCL result in similar performance, rather than NCCL being worse.
I'm not sure if it's related to vLLM's benchmark testing, but we've tested multiple times, and based on the current performance, for the vLLM v0 engine, using NCCL communication and enabling OVERLAP COMM doesn't seem to make much of a performance difference. I'm quite confused at the moment.
PS: For determining whether to use NCCL communication, we modified the input parameter in the with_tensor_transport() function in the vLLM source code to 'cpu' for testing.
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-03-18T12:15:23Z
body: @ruisearch42 discussions in #51305 has resolved. Please review again. Thanks.
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-03-19T06:58:55Z
body: > Thanks for the change and continued effort.
> 
> I will try but not sure will have a detailed review this week due to oncall tasks. Maybe your colleagues could help review/approve first? especially for the NPU part.

Yes, we can review it with NPU part first. Thanks.
author_association: CONTRIBUTOR
user: login: Bye-legumes
created_at: 2025-03-20T03:44:26Z
body: ## Further Investigation Needed

I believe this PR has potential performance issues compared to the official `torch_npu`, which require further investigation. Based on my test results in [PR #47658](https://github.com/ray-project/ray/pull/47658), there are noticeable differences in performance.

### Performance Results on vLLM (Model: Qwen2.5-14B)

| Total Cards | TP/FP Configuration | default (mp) | ray | ray + hccl |
|------------|---------------------|--------------|--------------|--------------|
| 2          | TP = 2, PP = 1      | 1627.300708  | 1585.444219  | 1512.801163  |
|            | TP = 1, PP = 2      | <2           | <2           | **2440.919**  |
| 4          | TP = 2, PP = 2      | <2           | <2           | 2007.215     |
|            | TP = 4, PP = 1      | 1658.146417  | 1682.88043   | 1707.24      |
|            | TP = 1, PP = 4      | <2           | <2           | **2715.516**  |

Given the concerns about performance and the reliance on third-party C code (e.g., HCCL), I don’t think this PR should be merged at this stage as i dont think ray can integrate these codes into build stage (**Ideally we should keep minor modification to add new device**). We need a more structured approach to integrating NPU support in Ray, ideally leveraging official packages rather than introducing extensive third-party dependencies.

## Related PRs for Context

- **Nov 2024**: Refactored NCCL to communicator channel to work with HCCL  
  [#48607](https://github.com/ray-project/ray/pull/48607)
- **Sep 2024**: Initial HCCL support, and I just want this PR as example for further other device reference on how to implementation this 
  [#47658](https://github.com/ray-project/ray/pull/47658)
- **Feb 2025**: HCCL support for Ray Collective, already deployed at ByteDance and other companies for RL workloads  
  [#50790](https://github.com/ray-project/ray/pull/50790)

- We need a high-level strategy for NPU support in Ray, considering whether external integrations or official Ascend support are required.
- The current approach relies heavily on custom C code. We should explore whether official packages can be used to minimize maintenance overhead.

### Unresolved Issues:
1. How can we ensure compatibility across different `CANN + torch_npu` versions? Should we provide prebuilt images?
2. How should Ray support NPU testing to enable future contributions?

**Until these issues are addressed, I think we should need further discussion and planning at current stage.**

author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-03-20T12:07:03Z
body: @Bye-legumes Yes, there are indeed some differences in the performance test results. Could you share the testing method? I need to rule out any discrepancies caused by different testing approaches.

Additionally, regarding the HCCL wrapper, I suggest using the same HCCL library as RCCL. Given that the interface wrapper for HCCL involves minimal code, creating a separate project would significantly increase maintenance costs. I recommend integrating this HCCL wrapper directly into the Ray project.
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-03-25T08:02:55Z
body: @liuxsh9 @Bye-legumes @noemotiovon  Good day. Could you please review this PR? Thanks.
This PR only includes the generalization of device handling in CG, and the HCCL part has been removed. If you want to check the historical HCCL implementation, please click [this link](https://github.com/hipudding/ray/blob/raybackup0326/python/ray/experimental/channel/hccl_group.py)
author_association: CONTRIBUTOR
user: login: liuxsh9
created_at: 2025-03-27T15:38:22Z
body: > @liuxsh9 @Bye-legumes @noemotiovon Good day. Could you please review this PR? Thanks. This PR only includes the generalization of device handling in CG, and the HCCL part has been removed. If you want to check the historical HCCL implementation, please click [this link](https://github.com/hipudding/ray/blob/raybackup0326/python/ray/experimental/channel/hccl_group.py)

The refactoring looks good to me! It seems that CG will be easily to support more accelerators.
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-03-28T06:07:54Z
body: @ruisearch42 Hello, this PR is now ready for review. According to the issue #51574 description, this PR only includes modifications for multi-backend support and has temporarily removed the HCCL-related logic. The HCCL logic will be implemented in a separate PR. My colleague has already completed the review and fixed the review comments.
Could you please review it again? Thanks.
author_association: CONTRIBUTOR
user: login: ruisearch42
created_at: 2025-03-28T15:08:52Z
body: > @ruisearch42 Hello, this PR is now ready for review. According to the issue #51574 description, this PR only includes modifications for multi-backend support and has temporarily removed the HCCL-related logic. The HCCL logic will be implemented in a separate PR. My colleague has already completed the review and fixed the review comments. Could you please review it again? Thanks.

Thanks. Will take a look!
author_association: CONTRIBUTOR
user: login: stephanie-wang
created_at: 2025-03-31T19:08:03Z
body: Thanks for this contribution! I'm still looking through the code, but regarding the API changes - can we make it so that you don't need to modify the Device enum to use a custom accelerator type? It would be better if we can only keep the NVIDIA backend in the main codebase and allow the user to provide a custom Communicator and AcceleratorRuntime. This would help maintainability of the project, as we also don't currently have a good way to test all different device types.
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-04-01T02:25:48Z
body: Thanks @stephanie-wang @ruisearch42. I will fix comment issue ASAP.
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-04-01T02:27:57Z
body: > Thanks for this contribution! I'm still looking through the code, but regarding the API changes - can we make it so that you don't need to modify the Device enum to use a custom accelerator type? It would be better if we can only keep the NVIDIA backend in the main codebase and allow the user to provide a custom Communicator and AcceleratorRuntime. This would help maintainability of the project, as we also don't currently have a good way to test all different device types.

Does this mean that the backend connects to new devices by registering Communicator and AcceleratorRuntime?

I have considered integrating via with_tensor_transport(custom_Communicator), but this requires specifying the list of participating actors when creating the Communicator, and still need provide custome Acclerator runtime functions.

However, I have a question to ask @stephanie-wang: where should the backend Communicator code be archived? Since this Communicator will only be used within Ray, would it be appropriate to place it in a separate folder (e.g., custom_communicators), so that user programs can import the required Communicator(or auto select)? The quality of these Communicators would be ensured by the backend vendors.

For Device class, we can set a unified custom backend name, such as customer_accelerator, and dynamically modify the device name when AcceleratorRuntime loads—just like PyTorch’s privateuse1 mechanism.

author_association: CONTRIBUTOR
user: login: stephanie-wang
created_at: 2025-04-06T22:01:49Z
body: > > Thanks for this contribution! I'm still looking through the code, but regarding the API changes - can we make it so that you don't need to modify the Device enum to use a custom accelerator type? It would be better if we can only keep the NVIDIA backend in the main codebase and allow the user to provide a custom Communicator and AcceleratorRuntime. This would help maintainability of the project, as we also don't currently have a good way to test all different device types.
> 
> Does this mean that the backend connects to new devices by registering Communicator and AcceleratorRuntime?

Yes, what I am imagining is to have some API call like `ray.register_accelerator_runtime("xpu", XpuAcceleratorRuntime, XpuCommunicator)`. So that way custom AcceleratorRuntimes can live outside of the Ray codebase.

> 
> I have considered integrating via with_tensor_transport(custom_Communicator), but this requires specifying the list of participating actors when creating the Communicator, and still need provide custome Acclerator runtime functions.
> 
> However, I have a question to ask @stephanie-wang: where should the backend Communicator code be archived? Since this Communicator will only be used within Ray, would it be appropriate to place it in a separate folder (e.g., custom_communicators), so that user programs can import the required Communicator(or auto select)? The quality of these Communicators would be ensured by the backend vendors.

Ideally I would like to support CPU and CUDA runtimes out of the box and other Communicators can live outside of the Ray codebase for now. The problem with other accelerators is that right now we don't have a way to test them in CI.

> 
> For Device class, we can set a unified custom backend name, such as customer_accelerator, and dynamically modify the device name when AcceleratorRuntime loads—just like PyTorch’s privateuse1 mechanism.

Hmm I'm not sure if I understand this idea. Is it an alternative to the `ray.register_accelerator_runtime` call? Can it support multiple AcceleratorRuntimes?

author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-04-07T01:08:17Z
body: Thanks @stephanie-wang, I will update this PR soon.

I will first write a document describing the overall work of the PR and the integration method for third-party backends. We’ll discuss based on the document first, and then move on to the code implementation.
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-04-08T02:39:57Z
body: @stephanie-wang @ruisearch42 Good day. Here's the [doc of this feature](https://docs.google.com/document/d/1X72U5fhd-0VoP4nuo4xJB3cqKKDJSgpYD1VzRK8LHk4/edit?usp=sharing), feel free to add any comments on it.
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-05-06T01:02:15Z
body: @ruisearch42 @stephanie-wang This PR is ready for review now. Could you please review this PR again? Thanks.
author_association: CONTRIBUTOR
user: login: ruisearch42
created_at: 2025-05-07T00:43:06Z
body: > @ruisearch42 @stephanie-wang This PR is ready for review now. Could you please review this PR again? Thanks.

Thanks for the efforts! will do!
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-05-07T06:36:25Z
body: > Thanks for the contribution! I like the new APIs. Let's check if we can reduce the driver API; eventually it would be good to use this as a single-controller API for Ray GPU objects too.

Sorry, I didn’t quite understand what you meant. Could you please explain this sentence again?
“Let’s check if we can reduce the driver API; eventually it would be good to use this as a single-controller API for Ray GPU objects too.”
author_association: CONTRIBUTOR
user: login: stephanie-wang
created_at: 2025-05-07T18:00:06Z
body: > > Thanks for the contribution! I like the new APIs. Let's check if we can reduce the driver API; eventually it would be good to use this as a single-controller API for Ray GPU objects too.
> 
> Sorry, I didn’t quite understand what you meant. Could you please explain this sentence again? “Let’s check if we can reduce the driver API; eventually it would be good to use this as a single-controller API for Ray GPU objects too.”

Ah this was just the comment about removing NotImplemented methods for the CommunicatorHandle class. Eventually we would like to use this API to program collective operations in general Ray programs, not just in compiled graphs (see https://github.com/ray-project/ray/issues/51173)
author_association: CONTRIBUTOR
user: login: Bye-legumes
created_at: 2025-05-12T18:31:37Z
body: I think currently looks great, while currently, the assumption is that all the device that support this should support torch stream, right? Is that possible that we list the APIs that to support all these features at least?
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-05-13T08:03:25Z
body: > I think currently looks great, while currently, the assumption is that all the device that support this should support torch stream, right? Is that possible that we list the APIs that to support all these features at least?

Yes, because currently CG only supports out-of-band transmission of torch.tensor, the accelerator must support the PyTorch interface—not just the torch stream, but also the interfaces related to device and event. I can add the backend capability requirements in the comments of accelerator_context.py.
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-05-15T01:11:47Z
body: > Thanks, this looks good! Please address the last few comments and then we can merge it.

Sorry for all the typos—I should have enabled local spell checking. Thank you for going through so many rounds of review and providing suggestions for improvement. I’ll fix these issues as soon as possible.
author_association: CONTRIBUTOR
user: login: stephanie-wang
created_at: 2025-05-21T00:16:52Z
body: Looks like there are just some lint errors to fix.
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-05-22T09:32:18Z
body: @stephanie-wang @ruisearch42  Hi, Could you please merge this PR? Thanks.
author_association: CONTRIBUTOR
user: login: daiping8
created_at: 2025-10-10T09:15:10Z
body: @hipudding Hello, this is a great piece of work that I can learn from. However, I have a few questions—would you mind helping me clarify them?

1. I couldn't find `hccl_group.py`. Was it removed from the commit because `hccl_group` and `nccl_group` are logically similar and contain a lot of duplicated code?
2. Is the current method for adding a new backend to CG still as described in your PR—implementing an `_XcclGroup` class similar to `_NcclGroup`, and registering it via `register_accelerator_context`?

Thank you for your time.
author_association: CONTRIBUTOR
user: login: hipudding
created_at: 2025-10-20T03:53:03Z
body: > 1. I couldn't find `hccl_group.py`. Was it removed from the commit because `hccl_group` and `nccl_group` are logically similar and contain a lot of duplicated code?

Yes, this part has removed and leave a register method to add your own backend. Currently, hccl backend can't add to Ray project because ray community has no Ascend servers and can't test it. I keep a backup of hccl and hccl_group, you can see https://github.com/ray-project/ray/commit/0482edc6dbbbe7d5834ea3d29c6bdb1d1b8c252c for an example.

> 2. Is the current method for adding a new backend to CG still as described in your PR—implementing an `_XcclGroup` class similar to `_NcclGroup`, and registering it via `register_accelerator_context`?

Yes, you're right, Use this register and enable your backend.


author_association: CONTRIBUTOR