[RFC] A pluggable PayloadTransfer interface for SimpleStorage

1. Background and motivation

SimpleStorage currently uses ZeroMQ (ZMQ) for PUT and GET traffic between Managers and StorageUnits. This RFC proposes support for additional communication backends such as UCX and NIXL, so deployments can choose a transport suited to their network and use the same SimpleStorage read and write logic.

2. Proposal and scope

Introduce an internal PayloadTransfer interface. Both ends of a connection run the selected backend. The backend owns the complete protocol for each logical PUT or GET. It decides how to establish a connection, deliver the request, obtain the data, exchange any memory descriptions and report completion. SimpleStorage supplies the local storage operations and determines whether a write has been accepted.

The Manager can submit a batch of byte ranges in a single call. Each range, called a span, describes contiguous bytes and retains their memory owner. When BatchMeta determines the layout of a GET result, the Manager can also supply the spans where that result should be available. This gives a backend the opportunity to receive into the final output and avoid an intermediate copy. When the layout is unknown, the backend can supply the response memory itself.

The initial scope is contiguous CPU tensor spans. SimpleStorage resolves references to shared fields, runs any configured storage-side data parsers, and constructs dense or nested results in the order required by the public APIs. Objects and other layouts can use serialization or packing. Allocation and registration strategies are implementation choices. Device memory and migration between StorageUnits are separate topics.

3. Interface and storage operations

The Manager submits requests. The StorageUnit starts the service and supplies two local operations for the backend to call. The signatures below describe the proposed internal API. Protocol framing and handshakes belong to the selected implementation.

class PayloadTransfer:
    def listen(self, address=None, *, load_data, store_data): ...  # Returns peer_info.
    async def connect(self, peer_info): ...
    async def request(self, connection, header, *, send=(), recv=()): ...
    async def close(self): ...

The business header describes the operation, sample indexes, fields and any layout information needed to interpret the data. send is the ordered list of source spans for the request. recv describes the destination of the response. request() returns (reply_header, received_spans). Each backend defines the wire representation of these arguments and how it matches concurrent replies to requests.

Manager supplies	Meaning on success
recv=()	The response has a header and no payload, as with a PUT acknowledgement.
recv=(dst0, dst1, ...)	The response parts fill those exact destinations. The agreed count, order and lengths describe logical spans. The backend can split or combine messages on the wire.
recv=None	The backend supplies memory and returns spans that keep the response data alive.

StorageUnit operation	Responsibility
load_data(get_header)	Look up the requested fields and samples, apply any configured read parser, and return the actual response layout and ordered source spans.
store_data(put_header, readable_spans)	Interpret the input, run any configured write parser, validate the result and accept it into storage. Return acceptance or a storage error.

The backend chooses when to invoke these operations. For instance, it may call load_data first to discover the size of a GET response, then arrange the destination. On PUT it may prepare a destination and start a read from the Manager before calling store_data with the completed input. The StorageUnit’s storage functions interpret field contents and enforce the storage rules.

4. Startup and peer discovery

A StorageUnit starts its PayloadTransfer instance and publishes the resulting peer_info through TQ’s bootstrap mechanism, which distributes service information to clients. TQ can pass this value unchanged. The backend interprets its contents, which may be a service address or additional connection metadata. The storage functions stay in the StorageUnit process.

# StorageUnit
peer_info = unit_transfer.listen(
    load_data=storage.load_data,
    store_data=storage.store_data,
)

# Manager, after obtaining this StorageUnit's peer_info from TQ bootstrap
peer = await transfer.connect(peer_info)

sequenceDiagram
    participant S as StorageUnit
    participant B as TQ bootstrap
    participant M as Manager
    S->>S: Start backend with local storage operations
    S->>B: Publish peer_info
    M->>B: Discover target StorageUnit
    B-->>M: peer_info
    M->>S: Connect using the selected backend
    S-->>M: Connection usable
    M->>S: Issue PUT or GET

A Manager can start later and connect only to the StorageUnits it needs. Setup takes place between the peers that will communicate. The selected backend defines the handshake and makes progress on its outstanding requests. When the service shuts down, close() drains or safely terminates outstanding transfers before releasing transport resources.

5. A batch from producer to consumer

Suppose a producer writes a batch containing one field, input_ids. Each sample holds a contiguous CPU tensor of token IDs. The sequences have the same length in this example. A training process later requests this field for a set of samples chosen by TQ. For clarity, all samples in the example route to one StorageUnit.

The following pseudocode runs inside the storage Managers. The producer’s Manager has the input batch data and its write metadata write_meta. The consumer’s Manager has its own read metadata read_meta, whose indexes are already in the consumer’s requested order. Each Manager obtains peer by connecting to the target StorageUnit. In the example, span(row) refers to the bytes of a row, and the two checks verify storage acceptance and agreement with the expected read layout.

# Producer's Manager
write_header = {
    "op": "PUT",
    "global_indexes": write_meta.global_indexes,
    "fields": ["input_ids"],
    "field_schema": extract_field_schema(data),
}
source = tuple(span(row) for row in data["input_ids"].unbind(0))
reply, _ = await transfer.request(peer, write_header, send=source)
check_put_ok(reply)

# Consumer's Manager
indexes = read_meta.global_indexes
field = read_meta.field_schema["input_ids"]
output = torch.empty((len(indexes), *field["shape"]), dtype=field["dtype"])
destinations = tuple(span(row) for row in output.unbind(0))
read_header = {"op": "GET", "global_indexes": indexes, "fields": ["input_ids"]}
reply, _ = await transfer.request(peer, read_header, recv=destinations)
check_response_layout(reply, read_meta)

For PUT, the StorageUnit adapter obtains the input bytes and calls store_data. For GET, it calls load_data to get the source spans in the requested order. On success, row i of output holds the tokens for indexes[i]. The backend may place the bytes directly or copy them into the supplied destinations. The Manager uses that layout to construct the dense or nested result expected by the caller.

For a batch spread across StorageUnits, the Manager routes each peer’s samples with their corresponding output spans. Variable-length tensors can use per-sample shapes from BatchMeta when available. If a parser determines the output shape at run time, the Manager uses recv=None and builds the result from the returned spans and actual schema.

6. Backend protocols

The backend selected for a connection carries its PUT/GET request signaling, payload and completion notifications. ZMQ carries this traffic only for the ZMQ backend. NIXL provides transfer operations through plugins. A NIXL adapter must provide its own request delivery and descriptor exchange and report an explicit error if the selected plugin or configuration cannot support the operation.

Backend	Example implementation
ZeroMQ (ZMQ)	A Manager DEALER socket and a StorageUnit ROUTER socket exchange multipart messages containing headers and payload frames. A GET can copy received frames into the caller’s specified destinations. ZMQ send semantics.
UCX Active Messages (AM)	The adapter registers a callback for a message type. A GET header can trigger load_data. For a large rendezvous PUT, the callback receives a descriptor, the adapter chooses a destination and calls ucp_am_recv_data_nbx. Once the input is readable it can call store_data. Replies can travel on the same AM channel. UCX AM API.
NIXL	The adapter defines how peers exchange requests and remote memory descriptions, submit data operations and notify one another of completion. It may use a transfer-bound notification when the selected plugin documents the required ordering.

One possible RMA protocol lets the StorageUnit adapter initiate the data operation. For PUT the Manager supplies readable source descriptions and the StorageUnit reads the data into local memory. For a known-layout GET the Manager supplies writable destinations, and the StorageUnit writes the selected data there. These descriptions grant access to specified byte ranges for the duration of the transfer.

sequenceDiagram
    participant M as Manager
    participant P as Manager adapter
    participant U as StorageUnit adapter
    participant S as Local storage
    Note over P,U: One possible RMA protocol
    rect rgb(238, 246, 255)
        M->>P: PUT with source spans
        P->>U: Request header and readable source descriptions
        U->>P: Issue READ into local destinations
        P-->>U: Sample data
        U->>S: store_data with readable input
        S-->>U: Write accepted
        U-->>P: PUT acknowledgement
        P-->>M: Request complete
        Note over M: Notify Controller when all units accept
    end
    rect rgb(240, 250, 242)
        M->>P: GET with prepared output spans
        P->>U: Requested fields, indexes and writable destinations
        U->>S: load_data for requested samples
        S-->>U: Actual layout and source spans
        U->>P: WRITE into output spans
        U->>U: Confirm remote data availability
        U-->>P: Response layout and completion
        P-->>M: Output ready in request order
    end

The RMA adapter in this example sends the GET completion after establishing that the remote bytes are available. With UCX’s UCP interface it can wait for an endpoint flush after the writes. With NIXL an independent notification needs a separate guarantee that the transfer has completed. Other protocols can negotiate a destination after discovering the source layout or use ordinary message delivery. Their exchanges and copies are part of the cost to measure. UCP completion semantics · NIXL Backend Guide.

7. Completion, errors and memory lifetime

Event	Guarantee
request() returns successfully	The transfer has finished accessing the request sources. All response bytes are readable in the agreed order and within the declared bounds. A layout mismatch is an error.
A PUT is acknowledged	The StorageUnit has accepted the write. Only after every StorageUnit involved has accepted its part does the Manager notify the Controller that the fields are ready.
A PUT reply is lost	If the backend cannot establish whether the write was accepted, it reports the outcome as unknown. A retry depends on the failure phase and whether replay is safe.
A request times out, is cancelled, or loses its connection	The adapter retains memory that outstanding work might access until the access finishes or is safely terminated.

Source bytes must stay stable while the transport may read them. Stored values, returned tensors and any surviving views must remain usable for their normal lifetimes, including across later requests. Callers use ordinary tensor lifetime management. Backends may use ordinary allocations, memory registration, temporary staging or pools. A prepared destination specifies where the caller can read the result on return and leaves the method of filling it to the backend.

8. Validation and performance

Run identical workloads through the public SimpleStorage APIs for every supported adapter. Verify field values and order across multiple StorageUnits, known and unknown result layouts, concurrent requests and a Manager joining during traffic. Retain earlier results and views across later operations. Exercise storage rejection, allocation failure, lost replies, cancellation and reconnection while transfers are in flight.

Verify that NIXL and UCX adapters can complete their supported operations with the SimpleStorage ZMQ request socket unavailable. An unsupported configuration must surface as an error. Count each logical PUT or GET once and its underlying protocol messages separately.

Measure setup and steady state separately. Include Controller publication in end-to-end PUT latency and final result construction in GET latency. Vary tensor size, number of spans, concurrency and allocation reuse. Compare useful throughput, latency distributions and CPU cost. Report copying, registrations and memory consumption where applicable.

9. Questions for review

1. Can the request arguments and the two local storage operations represent the SimpleStorage workloads an adapter needs to support?
2. Are the completion and failure guarantees precise enough to determine when an application can use or reuse a buffer?
3. Which workloads and acceptance criteria should guide enabling a new backend by default?