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 – “nccl” means that tensors will be passed via NCCL. “auto” (default) means that tensor transport will be automatically determined based on the sender and receiver, either through NCCL or host memory.

  • 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.