reference/storage

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

Key-value stores and UCXX

This page lists the tensordicts that keep their entries on a key-value server, and the classes and functions that stream tensordicts between processes over UCXX. For how to use them, see Large datasets and storage backends and Distributed tensordicts.

Key-value stores

A TensorDictStore keeps its entries on a server that speaks the Redis protocol, such as Redis or Dragonfly. The data is on the server, and each read gets it from the server. A LazyStackedTensorDictStore is a lazy stack of tensordicts on such a server, so that its entries can have different shapes. You need the redis package and a running server.

Warning

LazyStackedTensorDictStore is experimental: its API can change in a minor release without a deprecation cycle.

Both classes have the methods of TensorDictBase, listed by task in Methods by task. Their pages list the methods that they add, such as from_store, close, and clear_redis. STORE_BACKENDS is the type of their backend argument.

TensorDictStore(*[, backend, host, port, ...])	A TensorDict whose entries live on a Redis-compatible server.
LazyStackedTensorDictStore(*[, backend, ...])	A lazy stack of tensordicts whose entries live on a Redis-compatible server.

STORE_BACKENDS

alias of Literal['redis', 'dragonfly']

UCXX transport

Warning

TensorDictPipe, TensorDictServer, send_tensordict, and recv_tensordict are experimental: their API can change in a minor release without a deprecation cycle.

UCXX is the Python interface of UCX, a communication library that chooses the transport between two processes, such as shared memory, InfiniBand, or TCP. These classes and functions need the optional ucxx package.

• A TensorDictPipe is a channel between two processes. After the first tensordict, it sends only the data of the next tensordicts with the same structure. The receiver writes each tensordict into the same buffer, so clone a received tensordict if you need it after the next receive.

• A TensorDictServer accepts connections and returns one pipe for each client.

• send_tensordict() and recv_tensordict() send or receive one tensordict over a UCXX endpoint, with its metadata each time.

TensorDictBase.asend() and TensorDictBase.arecv() send and receive through a pipe.

Import these names from tensordict._ucxx. It is their only import path for now, although the module name starts with an underscore.

TensorDictPipe(endpoint, *[, consolidated, ...])	A persistent point-to-point channel for streaming TensorDicts over UCXX.
TensorDictServer(port, *[, consolidated])	A one-to-many listener that accepts connections and yields TensorDictPipe instances.

send_tensordict(endpoint, td, *[, consolidated])	Send a TensorDict over a UCXX endpoint (stateless, always sends metadata).
recv_tensordict(endpoint[, td, device])	Receive a TensorDict from a UCXX endpoint (stateless, always expects metadata).

See also

• Large datasets and storage backends: start a server, write a tensordict to it, and read it back.

• Containers on storage backends: which containers work on a store.

• Distributed tensordicts: send tensordicts between processes.

• Distributed and remote: the send and receive methods of every tensordict.