ray.dag.DAGNode.with_tensor_transport#

DAGNode.with_tensor_transport(transport: str | Communicator | None = 'auto', device: Literal['default', 'cpu', 'gpu', 'cuda'] = 'default', _static_shape: bool = False, _direct_return: bool = False)[source]#

Configure the torch tensor transport for this node.

Parameters:

  • transport – Specifies the tensor transport mechanism. - “accelerator”: Tensors are communicated using accelerator-specific backends (e.g., NCCL, XLA, or vendor-provided transport). This is the recommended option for most use cases, as it supports extensibility and future hardware backends. - “nccl”: Tensors are passed explicitly via NCCL. This option is kept for backwards compatibility and may be removed in the future. Use “accelerator” instead unless you have legacy requirements. - “shm”: Tensors are passed via host shared memory and gRPC. Typically used when accelerator-based transport is unavailable or not suitable. - “auto” (default): The system automatically selects the appropriate transport mechanism based on the sender and receiver, usually preferring accelerator-based transport when available.

  • device – The target device to use for the tensor transport. “default”: The tensor will maintain its original device placement from the sender “cpu”: The tensor will be explicitly moved to CPU device in the receiver “gpu” or “cuda”: The tensor will be explicitly moved to GPU device in the receiver

  • _static_shape – A hint indicating whether the shape(s) and dtype(s) of tensor(s) contained in this value always remain the same across different executions of the DAG. If this is True, the transport will be more efficient.

  • _direct_return – Whether the tensor is sent directly or inside of other data. If a “nccl” transport is used, this allows the sender and receiver to eliminate performance overhead from an additional data transfer.