Ray Direct Transport (RDT)#

Ray normally stores objects in its CPU-based object store, then copies and deserializes them when a Ray task or actor accesses them. For GPU data, this can cause unnecessary and expensive data transfers. For example, passing a CUDA torch.Tensor from one Ray task to another requires a copy from GPU to CPU memory, then back to GPU memory.

With Ray Direct Transport (RDT), Ray stores objects and passes them directly between Ray actors. RDT augments the familiar Ray ObjectRef API in the following ways:

Keeping GPU data in GPU memory until a transfer is necessary.
Avoiding expensive serialization and copies to and from the Ray object store.
Using efficient data transports, such as the Gloo or NCCL collective communication libraries or point-to-point remote direct memory access (RDMA) through NVIDIA’s NIXL, to transfer data directly between devices, including CPUs and GPUs.

Note

RDT is currently in alpha and doesn’t support all Ray Core APIs yet. Future releases might introduce breaking API changes. For more details, see the limitations section.

Getting started#

Tip

RDT currently supports torch.Tensor objects created by Ray actor tasks. Future releases might add support for other data types and for Ray non-actor tasks.

This walkthrough shows how to create and use RDT with different tensor transports, which are the mechanisms that transfer tensors between actors. RDT currently supports the following tensor transports:

Gloo: A collective communication library for PyTorch and CPUs.
NVIDIA NCCL: A collective communication library for NVIDIA GPUs.
NVIDIA NIXL: A library for accelerating point-to-point transfers through RDMA, especially between various types of memory and NVIDIA GPUs. NIXL runs on libfabric on AWS instances with Elastic Fabric Adapter (EFA), and on Unified Communication X (UCX) everywhere else.

To keep the walkthrough easy to follow, start with the Gloo transport, which works without any physical GPUs.

Usage with Gloo (CPUs only)#

Installation#

Note

Under construction.

Walkthrough#

To get started, define an actor class and a task that returns a torch.Tensor:

import torch
import ray


@ray.remote
class MyActor:
    def random_tensor(self):
        return torch.randn(1000, 1000)

As written, when the actor task returns the torch.Tensor, Ray copies it into its CPU-based object store. For CPU-based tensors, this can require an expensive step to copy and serialize the object. GPU-based tensors also require a copy to and from CPU memory.

To enable RDT, use the tensor_transport option in the @ray.method decorator.

@ray.remote
class MyActor:
    @ray.method(tensor_transport="gloo")
    def random_tensor(self):
        return torch.randn(1000, 1000)

You can add this decorator to any actor task that returns a torch.Tensor, or that returns torch.Tensors nested inside other Python objects. Adding this decorator changes Ray’s behavior in the following ways:

When returning the tensor, Ray stores a reference to the tensor instead of copying it to CPU memory.
When you pass the ray.ObjectRef to another task, Ray uses Gloo to transfer the tensor to the destination task.

For the second behavior to work, you only need to add the @ray.method(tensor_transport) decorator to the actor task that returns the tensor. Don’t add it to actor tasks that consume the tensor, unless those tasks also return tensors.

The second behavior also requires that you first create a collective group of actors.

Creating a collective group#

To create a collective group for use with RDT, do the following:

Create multiple Ray actors.
Create a collective group on the actors with the ray.experimental.collective.create_collective_group function. The backend you specify must match the tensor_transport in the @ray.method decorator.

The following example shows both steps:

import torch
import ray
from ray.experimental.collective import create_collective_group


@ray.remote
class MyActor:
    @ray.method(tensor_transport="gloo")
    def random_tensor(self):
        return torch.randn(1000, 1000)

    def sum(self, tensor: torch.Tensor):
        return torch.sum(tensor)


sender, receiver = MyActor.remote(), MyActor.remote()
# The tensor_transport specified here must match the one used in the @ray.method
# decorator.
group = create_collective_group([sender, receiver], backend="torch_gloo")

The actors can now communicate directly through Gloo. To destroy the group, use the ray.experimental.collective.destroy_collective_group function. After you call this function, you can create a new collective group on the same actors.

Passing objects to other actors#

With a collective group in place, you can create and pass RDT objects between the actors. The following example shows the full flow:

import torch
import ray
from ray.experimental.collective import create_collective_group


@ray.remote
class MyActor:
    @ray.method(tensor_transport="gloo")
    def random_tensor(self):
        return torch.randn(1000, 1000)

    def sum(self, tensor: torch.Tensor):
        return torch.sum(tensor)


sender, receiver = MyActor.remote(), MyActor.remote()
group = create_collective_group([sender, receiver], backend="torch_gloo")

# The tensor will be stored by the `sender` actor instead of in Ray's object
# store.
tensor = sender.random_tensor.remote()
result = receiver.sum.remote(tensor)
print(ray.get(result))

When you pass the ray.ObjectRef to another task, Ray uses Gloo to transfer the tensor directly from the source actor to the destination actor instead of through the default object store. The @ray.method(tensor_transport) decorator is only on the actor task that returns the tensor. Once you add this hint, the receiving actor task receiver.sum automatically uses Gloo to receive the tensor. In this example, because MyActor.sum doesn’t have the @ray.method(tensor_transport) decorator, it uses the default Ray object store transport to return torch.sum(tensor).

RDT also supports passing tensors nested inside Python data structures, and actor tasks that return multiple tensors, as in the following example:

import torch
import ray
from ray.experimental.collective import create_collective_group


@ray.remote
class MyActor:
    @ray.method(tensor_transport="gloo")
    def random_tensor_dict(self):
        return {"tensor1": torch.randn(1000, 1000), "tensor2": torch.randn(1000, 1000)}

    def sum(self, tensor_dict: dict):
        return torch.sum(tensor_dict["tensor1"]) + torch.sum(tensor_dict["tensor2"])


sender, receiver = MyActor.remote(), MyActor.remote()
group = create_collective_group([sender, receiver], backend="torch_gloo")

# Both tensor values in the dictionary will be stored by the `sender` actor
# instead of in Ray's object store.
tensor_dict = sender.random_tensor_dict.remote()
result = receiver.sum.remote(tensor_dict)
print(ray.get(result))

Passing RDT objects to the actor that produced them#

You can also pass RDT ray.ObjectRefs to the actor that produced them. This avoids any copies and provides a reference to the same torch.Tensor that the actor created earlier, as in the following example:

import torch
import ray
import pytest
from ray.experimental.collective import create_collective_group


@ray.remote
class MyActor:
    @ray.method(tensor_transport="gloo")
    def random_tensor(self):
        return torch.randn(1000, 1000)

    def sum(self, tensor: torch.Tensor):
        return torch.sum(tensor)


sender, receiver = MyActor.remote(), MyActor.remote()
group = create_collective_group([sender, receiver], backend="torch_gloo")

tensor = sender.random_tensor.remote()
# Pass the ObjectRef back to the actor that produced it. The tensor will be
# passed back to the same actor without copying.
sum1 = sender.sum.remote(tensor)
sum2 = receiver.sum.remote(tensor)
assert torch.allclose(*ray.get([sum1, sum2]))

Note

Ray only keeps a reference to the tensor that your code creates, so the tensor objects are mutable. If sender.sum modified the tensor in the preceding example, receiver.sum would also see the changes. This differs from the normal Ray Core API, which always makes an immutable copy of data that actors return.

`ray.get`#

You can also use the ray.get function as usual to retrieve the result of an RDT object. However, by default ray.get uses the same tensor transport as the one specified in the @ray.method decorator. For collective-based transports, this doesn’t work if the caller isn’t part of the collective group.

Therefore, specify the Ray object store as the tensor transport explicitly by setting _use_object_store in ray.get.

# Wrong example of ray.get(). Since the tensor transport in the @ray.method decorator is Gloo,
# ray.get() will try to use Gloo to fetch the tensor, which is not supported
# because the caller is not part of the collective group.
with pytest.raises(ValueError) as e:
    ray.get(tensor)

assert (
    "ray.get is not allowed on RDT objects using the two-sided transport"
    in str(e.value)
)

# Correct example of ray.get(), using the object store to fetch the RDT object because the caller
# is not part of the collective group.
print(ray.get(tensor, _use_object_store=True))
# torch.Tensor(...)

Object mutability#

Unlike objects in the Ray object store, RDT objects are mutable, meaning that Ray only holds a reference to the tensor and doesn’t copy it until a transfer is requested. If the actor that returns a tensor also keeps a reference to the tensor, and later modifies it in place while Ray is still storing the tensor reference, receiving actors might see some or all of the changes.

The following example shows what can go wrong:

import torch
import ray
from ray.experimental.collective import create_collective_group


@ray.remote
class MyActor:
    @ray.method(tensor_transport="gloo")
    def random_tensor(self):
        self.tensor = torch.randn(1000, 1000)
        # After this function returns, Ray and this actor will both hold a
        # reference to the same tensor.
        return self.tensor

    def increment_and_sum_stored_tensor(self):
        # NOTE: In-place update, while Ray still holds a reference to the same tensor.
        self.tensor += 1
        return torch.sum(self.tensor)

    def increment_and_sum(self, tensor: torch.Tensor):
        return torch.sum(tensor + 1)


sender, receiver = MyActor.remote(), MyActor.remote()
group = create_collective_group([sender, receiver], backend="torch_gloo")

tensor = sender.random_tensor.remote()
tensor1 = sender.increment_and_sum_stored_tensor.remote()
# Wait for sender.increment_and_sum_stored_tensor task to finish.
tensor1 = ray.get(tensor1)
# Receiver will now receive the updated value instead of the original.
tensor2 = receiver.increment_and_sum.remote(tensor)

try:
    # This assertion will fail because sender.increment_and_sum_stored_tensor
    # modified the tensor in place before sending it to
    # receiver.increment_and_sum.
    assert torch.allclose(tensor1, ray.get(tensor2))
except AssertionError:
    print("AssertionError: sender and receiver returned different sums.")

In this example, the sender actor returns a tensor to Ray, but it also keeps a reference to the tensor in its local state. Then, in sender.increment_and_sum_stored_tensor, the sender actor modifies the tensor in place while Ray is still holding the tensor reference. Then, the receiver.increment_and_sum task receives the modified tensor instead of the original, so the assertion fails.

To fix this kind of error, use the ray.experimental.wait_tensor_freed function to wait for Ray to release all references to the tensor, so that the actor can safely write to the tensor again. wait_tensor_freed unblocks once all tasks that depend on the tensor have finished executing and all corresponding ObjectRefs have gone out of scope. Ray identifies the tasks that depend on the tensor by tracking which tasks take the ObjectRef corresponding to the tensor as an argument.

Here’s a fixed version of the earlier example.

import torch
import ray
from ray.experimental.collective import create_collective_group


@ray.remote
class MyActor:
    @ray.method(tensor_transport="gloo")
    def random_tensor(self):
        self.tensor = torch.randn(1000, 1000)
        return self.tensor

    def increment_and_sum_stored_tensor(self):
        # 1. Sender actor waits for Ray to release all references to the tensor
        # before modifying the tensor in place.
        ray.experimental.wait_tensor_freed(self.tensor)
        # NOTE: In-place update, but Ray guarantees that it has already released
        # its references to this tensor.
        self.tensor += 1
        return torch.sum(self.tensor)

    def increment_and_sum(self, tensor: torch.Tensor):
        # Receiver task remains the same.
        return torch.sum(tensor + 1)


sender, receiver = MyActor.remote(), MyActor.remote()
group = create_collective_group([sender, receiver], backend="torch_gloo")

tensor = sender.random_tensor.remote()
tensor1 = sender.increment_and_sum_stored_tensor.remote()
# 2. Skip `ray.get`` because `wait_tensor_freed`` will block until all
# references to `tensor` are freed, so calling `ray.get` here would cause a
# deadlock.
# tensor1 = ray.get(tensor1)
tensor2 = receiver.increment_and_sum.remote(tensor)

# 3. Delete all references to `tensor`, to unblock wait_tensor_freed.
del tensor

# This assertion will now pass.
assert torch.allclose(ray.get(tensor1), ray.get(tensor2))

The main changes are the following:

sender calls wait_tensor_freed before modifying the tensor in place.
The driver skips ray.get because wait_tensor_freed blocks until all ObjectRefs pointing to the tensor are freed, so calling ray.get here would cause a deadlock.
The driver calls del tensor to release its reference to the tensor. Again, this is necessary because wait_tensor_freed blocks until all ObjectRefs pointing to the tensor are freed.

When you pass an RDT ObjectRef back to the same actor that produced it, Ray passes back a reference to the tensor instead of a copy, so the same kind of bug can occur. To help catch such cases, Ray prints a warning if you pass an RDT object to the actor that produced it and a different actor, as in the following example:

import torch
import ray
from ray.experimental.collective import create_collective_group


@ray.remote
class MyActor:
    @ray.method(tensor_transport="gloo")
    def random_tensor(self):
        return torch.randn(1000, 1000)

    def increment_and_sum(self, tensor: torch.Tensor):
        # In-place update.
        tensor += 1
        return torch.sum(tensor)


sender, receiver = MyActor.remote(), MyActor.remote()
group = create_collective_group([sender, receiver], backend="torch_gloo")

tensor = sender.random_tensor.remote()
tensor1 = sender.increment_and_sum.remote(tensor)
tensor2 = receiver.increment_and_sum.remote(tensor)
# A warning will be printed:
# UserWarning: GPU ObjectRef(...) is being passed back to the actor that created it Actor(MyActor, ...). Note that GPU objects are mutable. If the tensor is modified, Ray's internal copy will also be updated, and subsequent passes to other actors will receive the updated version instead of the original.

try:
    # This assertion may fail because the tensor returned by sender.random_tensor
    # is modified in-place by sender.increment_and_sum while being sent to
    # receiver.increment_and_sum.
    assert torch.allclose(ray.get(tensor1), ray.get(tensor2))
except AssertionError:
    print("AssertionError: sender and receiver returned different sums.")

Usage with NCCL (NVIDIA GPUs only)#

Switching RDT to a different tensor transport takes only a few lines of code change. The following code is the Gloo example, modified to use NVIDIA GPUs and the NCCL library for collective GPU communication.

import torch
import ray
from ray.experimental.collective import create_collective_group


@ray.remote(num_gpus=1)
class MyActor:
    @ray.method(tensor_transport="nccl")
    def random_tensor(self):
        return torch.randn(1000, 1000).cuda()

    def sum(self, tensor: torch.Tensor):
        return torch.sum(tensor)


sender, receiver = MyActor.remote(), MyActor.remote()
group = create_collective_group([sender, receiver], backend="nccl")

# The tensor will be stored by the `sender` actor instead of in Ray's object
# store.
tensor = sender.random_tensor.remote()
result = receiver.sum.remote(tensor)
ray.get(result)

The main code differences are the following:

The @ray.method uses tensor_transport="nccl" instead of tensor_transport="gloo".
The code passes backend="nccl" to the ray.experimental.collective.create_collective_group function, instead of backend="torch_gloo".
The code creates the tensor on the GPU with the .cuda() method.

Usage with NIXL (CPUs or NVIDIA GPUs)#

Installation#

First, install NIXL with pip install nixl. For maximum performance, run the install_gdrcopy.sh script, for example install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "12.8" "x64". For the available OS versions, see the NVIDIA GDRCopy downloads for CUDA 12.8.

When Ray selects the UCX backend, as NIXL backend selection describes, set the following UCX environment variables, either so that UCX chooses the right transport from all options or to set your preferred transport yourself. These variables apply only to the UCX backend and have no effect when Ray runs the LIBFABRIC backend.

# Example UCX configuration, adjust according to your environment
$ export UCX_TLS=all  # or specify specific transports like "rc,ud,sm,^cuda_ipc" ..etc
$ export UCX_NET_DEVICES=all  # or specify network devices like "mlx5_0:1,mlx5_1:1"

On the LIBFABRIC backend, use libfabric environment variables instead, such as FI_PROVIDER=efa to pin the provider and FI_LOG_LEVEL=Debug for diagnostics. See the NIXL libfabric plugin documentation for more configuration and troubleshooting information.

NIXL backend selection#

Ray automatically selects the NIXL transport backend based on the available hardware:

AWS instances with EFA: Ray detects EFA devices and validates that a realistic-size CUDA memory registration succeeds before it selects the LIBFABRIC backend. If validation fails, GPUDirect is usually misconfigured, for example because nvidia-peermem or dmabuf support is missing. To detect EFA both on bare hosts and inside containers, Ray checks for the EFA network device at /sys/class/net/efa* and for rdma-verbs devices under /sys/class/infiniband that are bound to the kernel efa driver. The host network device isn’t visible inside a pod, so the verbs check is what makes detection work under Kubernetes. Ordinary InfiniBand or RoCE NICs also expose verbs devices, so Ray confirms the efa driver binding to avoid treating them as EFA.
All other environments: Ray uses the UCX backend.

This selection requires no configuration. On AWS EFA instances, make sure you’ve run the EFA installer, which installs both the EFA driver and libfabric. If LIBFABRIC validation fails at startup, see the NIXL libfabric plugin documentation for troubleshooting.

Walkthrough#

NIXL can transfer data between different devices, including CPUs and NVIDIA GPUs, but doesn’t require you to create a collective group ahead of time. Any actor that has NIXL installed in its environment can create and pass an RDT object.

Otherwise, usage is the same as in the Gloo example.

The following example uses NIXL to transfer an RDT object between two actors:

import torch
import ray


@ray.remote(num_gpus=1)
class MyActor:
    @ray.method(tensor_transport="nixl")
    def random_tensor(self):
        return torch.randn(1000, 1000).cuda()

    def sum(self, tensor: torch.Tensor):
        return torch.sum(tensor)

    def produce(self, tensors):
        refs = []
        for t in tensors:
            refs.append(ray.put(t, _tensor_transport="nixl"))
        return refs

    def consume_with_nixl(self, refs):
        # ray.get will also use NIXL to retrieve the
        # result.
        tensors = [ray.get(ref) for ref in refs]
        sum = 0
        for t in tensors:
            assert t.device.type == "cuda"
            sum += t.sum().item()
        return sum


# No collective group is needed. The two actors just need to have NIXL
# installed.
sender, receiver = MyActor.remote(), MyActor.remote()

# The tensor will be stored by the `sender` actor instead of in Ray's object
# store.
tensor = sender.random_tensor.remote()
result = receiver.sum.remote(tensor)
ray.get(result)

Compared to the Gloo example, the main code differences are:

The @ray.method uses tensor_transport="nixl" instead of tensor_transport="gloo".
No collective group is needed.

ray.put and ray.get with NIXL#

Unlike with the collective-based tensor transports, Gloo and NCCL, the ray.get function can use NIXL to retrieve a copy of the result. By default, ray.get uses the tensor transport specified in the @ray.method decorator.

# ray.get will also use NIXL to retrieve the
# result.
print(ray.get(tensor))
# torch.Tensor(...)

You can also use NIXL to retrieve the result from references created by ray.put.

tensor1 = torch.randn(1000, 1000).cuda()
tensor2 = torch.randn(1000, 1000).cuda()
refs = sender.produce.remote([tensor1, tensor2])
ref1 = receiver.consume_with_nixl.remote(refs)
print(ray.get(ref1))

Summary#

With RDT, Ray stores objects and passes them directly between Ray actors, using accelerated transports such as Gloo, NCCL, and NIXL. Keep the following main points in mind:

If you use a collective-based tensor transport, Gloo or NCCL, you must create a collective group ahead of time. NIXL only requires all involved actors to have NIXL installed.
Unlike objects in the Ray object store, RDT objects are mutable, meaning that Ray only holds a reference to the stored tensors, not a copy.
Otherwise, you can use actors as normal.

For a full list of limitations, see the limitations section.

Microbenchmarks#

Note

Under construction.

Limitations#

RDT is currently in alpha and has the following limitations, which future releases might address:

Support for torch.Tensor objects only.
Support for Ray actors only, not Ray tasks.
Support for the following transports: Gloo, NCCL, and NIXL.
Support for CPUs and NVIDIA GPUs only.
RDT objects are mutable. Ray only holds a reference to the tensor and doesn’t copy it until a transfer is requested. If the application code also keeps a reference to a tensor before returning it and modifies the tensor in place, the receiving actor might see some or all of the changes.
await on an RDT object ref is temporarily unsupported.

The collective-based, two-sided tensor transports, Gloo and NCCL, have the following limitations:

Only the process that created the collective group can submit actor tasks that return and pass RDT objects. If the creating process passes the actor handles to other processes, those processes can submit actor tasks as usual, but can’t use RDT objects.
Similarly, the process that created the collective group can’t serialize and pass RDT ray.ObjectRefs to other Ray tasks or actors. Instead, you can only pass the ray.ObjectRefs as direct arguments to other actor tasks, and those actors must be in the same collective group.
Each actor can only be in one collective group per tensor transport at a time.
No support for ray.put.
No support for out-of-order actors such as async actors or actors with max_concurrency greater than 1.

Because of a known issue, Ray doesn’t currently support storing different GPU objects at the same actor with NIXL when the objects contain an overlapping but not equal set of tensors. To support this pattern, make sure the first ObjectRef has gone out of scope before you store the same tensors again in a second object.

@ray.remote(num_gpus=1)
class Actor:
    def __init__(self):
        self.tensor1 = torch.tensor([1, 2, 3])
        self.tensor2 = torch.tensor([4, 5, 6])
        self.tensor3 = torch.tensor([7, 8, 9])

    @ray.method(tensor_transport="nixl")
    def send_dict1(self):
        return {"round1-1": self.tensor1, "round1-2": self.tensor2}

    @ray.method(tensor_transport="nixl")
    def send_dict2(self):
        return {"round2-1": self.tensor1, "round2-3": self.tensor3}

    def sum_dict(self, dict):
        return sum(v.sum().item() for v in dict.values())


sender, receiver = Actor.remote(), Actor.remote()
ref1 = sender.send_dict1.remote()
result1 = receiver.sum_dict.remote(ref1)
print(ray.get(result1))
ref2 = sender.send_dict2.remote()
result2 = receiver.sum_dict.remote(ref2)
try:
    print(ray.get(result2))
except ValueError as e:
    print("Error caught:", e)

Error handling#

Application-level errors, which are exceptions that your code raises, don’t destroy the collective group. Instead, Ray propagates them to any dependent tasks, as for non-RDT Ray objects.
If a system-level error occurs during a Gloo or NCCL collective operation, the collective group is destroyed and the actors are killed to prevent any hanging.
If a system-level error occurs during a NIXL transfer, Ray or NIXL stops the transfer with an exception, and Ray raises the exception in the dependent task or in the ray.get call on the NIXL object ref.
System-level errors include the following:
- Errors internal to the third-party transport, for example NCCL network errors.
- Actor or node failures.
- Transport errors from tensor device and transport mismatches, for example a CPU tensor when using NCCL.
- Ray RDT object fetch timeouts. To override the timeout, set the RAY_rdt_fetch_fail_timeout_milliseconds environment variable.
- Any unexpected system bugs.

Advanced: Registering a custom tensor transport#

You can register custom tensor transports at runtime for use with RDT. To add a custom tensor transport, implement the abstract interface ray.experimental.TensorTransportManager and register it with ray.experimental.register_tensor_transport.

For a complete guide on implementing custom tensor transports, including detailed documentation of all required methods, see Implementing a custom tensor transport (Advanced).

Advanced: RDT internals#

Note

Under construction.

Table of contents#

For more details about Ray Direct Transport, see the following pages:

Implementing a custom tensor transport (Advanced)

Ray Direct Transport (RDT)#

Getting started#

Usage with Gloo (CPUs only)#

Installation#

Walkthrough#

Creating a collective group#

Passing objects to other actors#

Passing RDT objects to the actor that produced them#

ray.get#

Object mutability#

Usage with NCCL (NVIDIA GPUs only)#

Usage with NIXL (CPUs or NVIDIA GPUs)#

Installation#

NIXL backend selection#

Walkthrough#

ray.put and ray.get with NIXL#

Summary#

Microbenchmarks#

Limitations#

Error handling#

Advanced: Registering a custom tensor transport#

Advanced: RDT internals#

Table of contents#

`ray.get`#