set_target_device_for_ref#

ray.experimental.set_target_device_for_ref(ref: ray._raylet.ObjectRef, target_device: str)[source]#

Set the target device for an RDT ObjectRef to fetch tensors onto when ray.get is called.

Ray allocates the receive buffers on target_device instead of the source device. This enables cross-device transfers (e.g., fetching a CUDA tensor onto CPU or vice versa) for transports that support it (e.g., NIXL). If the transport does not support this feature, an exception will be raised during ray.get.

This overwrites any target buffer or target device previously set for the ref.

Parameters:

  • ref – The ObjectRef to set the target for. The ref must be for an RDT object.

  • target_device – The device to fetch the tensors onto, as a device string (e.g. "cpu", "cuda:0"). This may differ from the source device to perform a cross-device transfer.

PublicAPI (alpha): This API is in alpha and may change before becoming stable.