[WIP][feat] optimize NIXL-UCX payload transfer for SimpleStorage

[Ginray]

Motivation

PR #163 introduced the PayloadTransfer abstraction and the optional NIXL-UCX Host-to-Host payload backend for SimpleStorage.

This PR further optimizes and hardens the NIXL-UCX path for RL workloads. The main goals are to reduce avoidable Host-side copies, support large frame-native payloads, avoid blocking SimpleStorage workers during GET transfers, and make buffer / transfer lifecycle handling safer under concurrency and failure.

This work follows the optimization plan outlined in the Payload Transfer roadmap.

Summary

• Replace the packed NIXL payload format with frame-native scatter WRITE, so encoded frames can be transferred directly without rebuilding a payload-sized packed buffer.
• Cache and reuse registered receive buffers after their leases are released. A lease keeps a buffer unavailable for reuse while decoded tensors or views still reference its memory; idle cached buffers are bounded by receive_buffer_cache_mb per runtime.
• Reuse leased receive MRs as send sources when writable Host frames already reference them; other writable contiguous Host frames can be registered directly without an extra staging copy.
• Add deferred GET responses, allowing the SimpleStorage worker to continue processing requests while the NIXL transfer completes.
• Use a single-worker executor per remote peer: transfers to the same peer remain serialized, while different peers can progress concurrently.
• Harden failure and shutdown handling with receive-buffer quarantine, retained send resources, failed-peer isolation, and ordered runtime teardown.

Design

ZMQ remains the control plane, while non-empty Host payloads use NIXL-UCX.

SimpleStorage
     |
     | prepare / commit / response
     v
    ZMQ
     |
     +------------------------------+
                                    |
encoded frames                      |
     |                              |
     v                              v
NIXL-UCX scatter WRITE      registered receive buffer
                                    |
                                    v
                           zero-copy frame views
                                    |
                                    v
                                  decode

The implementation intentionally keeps the scope limited to the SimpleStorage Host payload fast path rather than introducing a general-purpose NIXL transport or MR management subsystem.

Compatibility

• Existing SimpleStorage public APIs remain unchanged.
• ZMQ remains the default payload backend.
• NIXL-UCX remains opt-in.
• Existing PUT / GET / CLEAR semantics are unchanged.
• NIXL transfer failures are reported explicitly and do not silently fall back to ZMQ.

Validation

Validation was performed on the current PR head with both two-node NIXL-UCX tests and local RL integration tests.

• Two-node tests verified the payload path over rma(rc_verbs/hns_0:1), including PUT / GET / CLEAR, datatype and nested-object handling, receive-buffer lease reuse, same-peer / different-peer concurrency, and controlled peer failure.
• PUT / GET correctness passed for payloads from 1 MiB through single-frame 5 GiB and multi-frame 8 GiB (3+3+2 GiB). No 32-bit packed-length truncation or StorageUnit splitting was observed.
• Source instrumentation confirmed that contiguous Host tensors were registered at their original data_ptr() addresses, with matching GET contents and SHA-256 hashes. Existing receive MRs were reused as send sources without another registration. No payload-sized staging copy was observed; small metadata frames were still copied.
• Receive-cache checks passed with 0, 4, and 256 MiB budgets. Idle buffers stayed within the configured budget. A GET already holding its source retained the original data across overwrite or CLEAR; subsequent GETs returned the new value or reported the key missing.
• In the local concurrency check, a healthy peer completed in about 0.268 s while another peer was delayed by 1.5 s. Concurrent GETs, probes, and CLEAR requests returned the expected responses.
• PUT and GET retried after a lost READY only when a matching CANCEL acknowledgement was received. A lost PUT COMMIT reply produced exactly 1 PREPARE, 1 COMMIT, and 0 CANCEL requests; the client reported an unknown outcome without replay, and a later GET confirmed the stored data. GET completion failures quarantined the target buffer.
• Lifecycle checks covered repeated close, pending receives, leased and quarantined buffers, and close during WRITE. Killing the receiver during a 512 MiB WRITE returned an explicit error; the manager made one attempt and did not replay or send CANCEL.
• All 1,000 PUT/GET/CLEAR cycles against a fixed healthy RDMA peer passed. The client and StorageUnit each registered one receive buffer and reused it 999 times. QP count stayed at 2, FD count remained stable, and sampled pending, leased, quarantined, and outstanding counts were zero after CLEAR.
• VIME Bridge RDMA checks matched rollout fields and hashes. Local VIME training proceeded normally through 10/10 steps, with 10 PUTs, 10 trainer fetches, and 10 CLEARs.

Follow-up

1. Continue reviewing whether the current lifecycle and failure-handling design can be simplified without weakening correctness.

2. Further validate performance and memory behavior in larger-scale distributed RL workloads.

3. Consider adding a resource limit to the idle registered receive pool. This may depend on the results of integration with RL frameworks, and I have not yet decided on the right approach.
   Current status: implemented a bounded idle receive-buffer cache with receive_buffer_cache_mb (default: 256 MiB per runtime); oversized transient buffers are released after their leases end.

4. Connect the runtime diagnostics to the existing TransferQueue metrics / Prometheus path.
   Current status: payload runtime diagnostics are exported through the existing metrics path as tq_storage_payload_* metrics.

[ascend-robot]

CLA Signature Pass

Ginray, thanks for your pull request. All authors of the commits have signed the CLA. 👍

[ascend-robot]

CLA Signature Pass

Ginray, thanks for your pull request. All authors of the commits have signed the CLA. 👍

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	⚠️ Failed 2026-10-12T12:41:12.320467Z	90ad2c4	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[0oshowero0]

[AI Review]

Summary. The direction is good: replacing the packed format with frame-native scatter WRITE is a net simplification. In its current form, though, I'd suggest addressing two issues before merging: a performance regression on the default ZMQ backend, and a failure mode where one uncertain transfer permanently disables a client ↔ storage-unit path with no recovery owner. A few reliability risks also need validation under realistic RL patterns, and the change would be much easier to review if split.

Locally, python -m compileall -q transfer_queue tutorial tests passes and 219 tests pass (test_nixl_payload_lifecycle, test_payload_transfer, test_metrics, test_simple_storage_unit, test_storage_request_retry, test_async_simple_storage_manager, test_zmq_socket_pool). NIXL was not installed, so the NIXL tests ran against _FakeAgent only.

Scope note. Against rdma, this PR also contains a merge of main (332cc69), hence +8380/−765 across 47 files; the feature itself is 90ad2c4 (12 files). The merge carries non-trivial conflict resolution (mapping NIXL timeouts to PayloadTransferTimeout for the #171 retry path, close() ordering, worker dispatch). It would be good to land the main → rdma sync as its own PR so that resolution gets reviewed on its own.

Should fix before merge

1. Default ZMQ backend: GET responses are serialized twice.
_process_one_worker_request serializes the response to measure wire size, returns the message, and the worker loop serializes it again:

TransferQueue/transfer_queue/storage/simple_storage.py

Lines 1187 to 1198 in 90ad2c4

	response_frames = response_msg.serialize()
	if operation == ZMQRequestType.GET_DATA: # type: ignore[arg-type]
	# This end serializes the get response, so its frames give the true wire size.
	log_heavy_operation(
	self.storage_unit_id,
	"get",
	time.perf_counter() - started,
	sum(frame_nbytes(frame) or 0 for frame in response_frames),
	f"samples={len(request_msg.body.get('global_indexes', []))} "
	f"fields={list(request_msg.body.get('fields', []))}",
	)
	return identity, response_msg

TransferQueue/transfer_queue/storage/simple_storage.py

Lines 1049 to 1050 in 90ad2c4

	else:
	worker_socket.send_multipart([identity] + response_msg.serialize(), copy=False)

For GET this walks the whole payload twice on the single storage worker thread (non-contiguous tensors are copied by .contiguous(), non-tensor values are pickled). A 512-sample response took ~3 ms per serialize() locally, scaling with payload size. This contradicts "ZMQ behavior unchanged". Returning the serialized frames from the helper fixes it.

2. One uncertain failure permanently disables a path, with no recovery.

TransferQueue/transfer_queue/storage/payload_transfer/nixl.py

Lines 343 to 347 in 90ad2c4

	if safe_to_release:
	self.cancel_receive(transfer_id)
	else:
	self._failed_get_targets.add(target_id)
	self.quarantine_receive(transfer_id)

• Client: a GET that times out or is cancelled after COMMIT adds the storage unit to _failed_get_targets; every later GET to it fails.
• Sender runtime: an uncertain WRITE adds the peer to _failed_peers; every later send to it fails.
• The docs say to "recreate the runtime through the owning worker's error handling", but nothing in TQ recreates payload_transfer: AsyncSimpleStorageManager and SimpleStorageUnit create it once. The [fix] Retry a lost storage-unit request instead of failing the job #171 retry does not apply either, since this path raises PayloadTransferError.
• This is reachable: the storage worker is single-threaded, so if it is held before handling COMMIT (a large GET encode, a checkpoint save) and the WRITE is slow, the client's 200 s ZMQ timeout can fire first. In a long RL job, one hiccup leaves that client unable to read from that storage unit for the rest of the run.

For memory safety, quarantining the one receive buffer is already sufficient: it is never reused, and later GETs use a fresh transfer_id and buffer, so a late WRITE can only land in the quarantined memory. GET is read-only, so it can also be retried after quarantine. Suggest dropping the per-target permanent failure in favor of a byte budget on quarantined memory (fail fast only when exceeded). On the sender side, either reset the remote agent or fail the job explicitly rather than leaving a half-working client.

Reliability risks to validate

3. On storage units, all stored data becomes registered (pinned) memory, retained per whole PUT.
After PUT commit, decoded tensors are views into the receive buffer, and put_data stores them without copying:

TransferQueue/transfer_queue/storage/simple_storage.py

Lines 276 to 277 in 90ad2c4

	for key, val in zip(global_indexes, values, strict=True):
	field_dict[key] = val

• receive_buffer_cache_mb only bounds idle buffers; pinned memory on a storage unit equals all stored data, subject to ulimit -l and NIC MR limits.
• If any small value from a PUT stays in memory, the whole buffer stays alive, including the bytes of large values already offloaded to SSD, so SSD offload ([feat] simple storage support ssd offload #162) saves little under NIXL. Partial clears behave the same way.
• RL usually accumulates several PUTs before CLEAR, so the storage-side cache rarely hits and each PUT registers a new buffer.

The 1000-cycle validation runs PUT/GET/CLEAR in lockstep, which is the one pattern that hides this. (The previous path copied into one per-PUT bytearray, so retention granularity was the same, but the memory was not pinned.) Consider keeping zero-copy on the client GET path only and copying out on storage-side PUT commit, which also makes the storage cache effective. At minimum, document this and validate with many outstanding PUTs.

4. The sender may remove and re-add the remote agent on almost every PUT.
Partial metadata is generated per registration, and the sender reloads whenever it differs from the previous one:

TransferQueue/transfer_queue/storage/payload_transfer/nixl_ucx_runtime.py

Lines 262 to 264 in 90ad2c4

	metadata = self._agent.get_partial_agent_metadata(
	scratch.registration, inc_conn_info=True, backends=["UCX"]
	)

TransferQueue/transfer_queue/storage/payload_transfer/nixl_ucx_runtime.py

Lines 646 to 652 in 90ad2c4

	def _load_remote_metadata(self, remote_name: str, metadata: bytes) -> None:
	previous = self._remote_metadata.get(remote_name)
	if previous == metadata:
	return
	if previous is not None:
	self._agent.remove_remote_agent(remote_name)
	loaded_name = self._agent.add_remote_agent(metadata)

With storage-side buffers held long-term (point 3), nearly every PUT presents new metadata. Whether this rebuilds UCX endpoints/QPs depends on NIXL, but "QP count stayed at 2" was measured with one buffer reused 999 times. Please check QP/FD counts with many un-cleared PUTs. If NIXL supports loading partial metadata incrementally for a known agent, the remove_remote_agent call may be unnecessary.

5. O(n²) overlap scan under the global runtime lock.

TransferQueue/transfer_queue/storage/payload_transfer/nixl_ucx_runtime.py

Lines 516 to 527 in 90ad2c4

	for scratch in self._leased_receive_buffers.values():
	if scratch.address <= address and address + size <= scratch.address + scratch.capacity:
	return _FrameSource(view, None, address, size)
	# NIXL resolves registrations by address; overlapping sources
	# need separate backing memory to release them independently.
	overlaps = any(
	address < source.address + source.size and source.address < address + size
	for source in self._registered_sources.values()
	)
	if not overlaps:
	return self._register_frame_source(view)

Jagged nested tensors are unbound into one frame per sample on PUT. Timing this scan alone: 0.22 s for 4096 frames and 0.86 s for 8192, all while holding the runtime RLock. Large register_memory calls (sources and receive buffers) also run under that lock, and the client calls prepare_receive on the asyncio event-loop thread, so the loop blocks too. "Different peers progress concurrently" only holds during WRITE polling. Coalescing small frames into one staging buffer and registering only large frames would remove both the quadratic scan and thousands of MR registrations.

6. Pending PUT/GET entries never expire.
Errors say "unknown or expired transfer_id", but nothing expires them. If a client dies between PREPARE and COMMIT, _PendingGet.frames keeps stored data alive, including data that is later CLEARed. This has existed since #163 (where deadline-based cleanup was deferred), but with leases the leaked memory is now pinned.

Simplicity and maintainability

• One commit bundles six changes: frame-native scatter, receive leases with a bounded cache, source-registration reuse, deferred GET with per-peer executors, failure isolation, and 17 metrics. Suggest splitting into at least: (a) drop the packed format (pure deletion, easy to approve), (b) leases, cache and source reuse, (c) deferred GET and per-peer executors, (d) failure handling, (e) metrics.
• NixlRuntime grew to ~12 containers/flags and 16 counters. Six diagnostics fields are never exported and only read by tests (idle/pending/leased_receive_buffers, receive_working_set_hwm_count/bytes, receive_buffer_acquisitions); they can go along with their bookkeeping (_active_receive_count/_bytes, etc.).
• _send_scatter tracks "safe to release" vs "uncertain" with four state variables (peer_state_uncertain flips three times, plus write_started, handle, error_message) and two cleanup blocks. Structuring it by phase (prepare, where failures release; post WRITE, uncertain afterwards; poll) would make it readable in one pass.
• The deferred-response plumbing uses two module-level helpers with 5 and 4 parameters, plus an ExitStack timing map threaded into _process_one_worker_request. If _drain_deferred_responses raises, the worker breaks: the worker thread exits while the storage unit still looks alive, with only a warning logged.
• serial_utils.initialize_packed_frame_table is now dead code (only its own unit test uses it).

Interfaces

• The handle_request -> ZMQMessage | DeferredResponse | None seam is reasonable, since the worker owns the socket. But PayloadTransfer.diagnostics may return any key while the exporter hard-codes 17 NIXL names, so unknown keys are silently dropped. Define them in one place.
• receiver_endpoint in GET COMMIT is redundant: the client builds it from the token and the storage unit checks it against the token (nixl.py#L304-L306). It can be dropped from the protocol.
• The one-line forwards on NixlPayloadTransfer (prepare_receive/send/receive/cancel_receive, plus the new quarantine_receive) are thin wrappers, and cast(int, ...) in factory.py has no runtime effect.

Test gaps

• Deferred responses are tested only through the two helpers; there is no SimpleStorageUnit-level test that interleaves a deferred GET with other requests.
• No test would catch the double serialization in point 1.
• _FakeAgent cannot validate the partial-metadata behavior in point 4.

[0oshowero0]

[AI Review] Follow-up: how common RDMA stacks manage registered memory

This expands on point 3 above. Common RDMA stacks handle registered memory in roughly three ways, and none of them registers a fresh buffer per request and then lets the lifetime of stored data decide when that buffer is freed.

1. Communication libraries (MPI, UCX, NCCL, TensorPipe, brpc): library-owned bounce buffers, user memory for large messages.
Small messages are copied through a fixed set of pre-registered buffers. Large messages use a rendezvous protocol directly on the user's memory: the library registers it through a registration cache (UCX rcache, Open MPI leave_pinned) and invalidates entries through malloc/munmap hooks. Users can also pre-register explicitly (ucp_mem_map, ncclCommRegister / ncclMemAlloc); otherwise NCCL copies through its internal registered buffers. The memory stays owned by the user; the library only caches registrations.

2. Transfer engines (NIXL, Mooncake Transfer Engine, 3FS USRBIO): the application provides registered memory.
Transfers only run between regions the application registered (register_memory, registerLocalMemory, hf3fs_iovcreate). vLLM, SGLang and Dynamo register the whole KV-cache pool once at startup and allocate blocks from it. Mooncake Store offers both forms: zero-copy get_into / put_from on buffers registered via register_buffer, and plain get / put through an internal pre-registered buffer plus a copy.

3. RDMA storage (Mooncake Store, FaRM, 3FS): one fixed pool registered at startup, allocated per object.
A PUT allocates a slice, the client RDMA-WRITEs into it, then commits; deleting the object frees its slice. Keeping all stored data in registered memory is normal here, under three conditions: registration happens once, the total equals the configured capacity, and memory is freed per object. brpc's RDMA IOBuf blocks and DPDK mbufs also hand applications refcounted views of registered memory, but from a fixed pool in small fixed-size blocks, with backpressure or a copy when the pool runs out.

Where this PR differs. On storage units, all three conditions are inverted:

• Every PUT allocates a new bytearray and calls register_memory, whose cost grows with size. [optim] MooncakeStore read path registers RDMA buffers on every get(), dominating end-to-end time #169 measured registration dominating end-to-end time on the MooncakeStore read path.
• Retention is per whole PUT: all fields and samples share one contiguous buffer, StorageUnitData keeps torch views (samples from unbind() share storage), and a single live view keeps the whole buffer alive.
• Pinned memory has no total bound: receive_buffer_cache_mb only covers idle buffers, while pinned memory equals all live data.

The client GET lease is milder, since consumers usually release a batch quickly, but a caller who keeps one small slice still pins the whole GET buffer.

There is already a precedent in this repo. The MooncakeStore read path (#175) follows the first pattern: a bounded pre-registered pool, one lease per read, an immediate copy into the caller's tensors, then release. The lease ends with the transfer, not with the data:

TransferQueue/transfer_queue/storage/clients/mooncake_client.py

Lines 583 to 598 in 90ad2c4

	lease = self._acquire_buffer(total) if total <= self._pool_share_bytes else None
	if lease is None:
	return group
	keys = [batch_keys[i] for i in group]
	try:
	self._batch_get_into_with_retry(keys, [lease.ptr + off for off in offsets], nbytes)
	_copy_buffer_content_into_tensors(lease, [batch_tensors[i] for i in group], offsets)
	except Exception:
	# A failed read can leave the buffer unreturnable; keep the original cause.
	try:
	lease.release()
	except Exception as release_error:
	logger.warning(f"Returning the pooled receive buffer failed: {release_error}")
	raise
	lease.release()

Suggestion.

• For the SimpleStorage NIXL path, use bounce buffers with a copy-out: at storage-side PUT commit, copy into ordinary memory and return the buffer to the cache right away; do the same on client GET, or keep zero-copy there with an explicit bound. Pinned memory then stays at cache plus in-flight transfers, the cache actually hits, and the number of registrations stays near constant. This also removes the SSD-offload interaction and the per-PUT metadata reload (points 3 and 4 above). The cost is one memcpy on the storage side, which the pre-PR path already paid; on 200/400 Gbps NICs a single-threaded copy may not keep up, so it should be measured and copied in parallel chunks if needed.
• If storage-side zero-copy is truly required, the right shape is the third pattern: a fixed pool registered at startup, slices allocated per sample or field, and freed on CLEAR. That amounts to reimplementing Mooncake Store inside SimpleStorage. Since TQ already ships a MooncakeStore backend, that need is better served there. The PR description also says it does not intend to introduce an MR management subsystem, and leases + cache + quarantine are already halfway to one.

[shun001]

Thanks for the detailed review. I have pushed the follow-up fixes for the two pre-merge findings and several related lifecycle issues.

• Default ZMQ GET: The storage worker now sends the frames already serialized by request processing. A GET response is serialized once rather than twice.
• Uncertain NIXL GET and unsafe native WRITE: An uncertain GET quarantines only that transfer's receive buffer; a later GET uses a new transfer ID and buffer. Quarantined memory has a configurable fail-fast budget. If a native sender enters an unsafe state, the entire payload-transfer runtime becomes terminal, and subsequent PUT/GET operations fail explicitly across peers. The owning job must propagate that failure and recreate its worker; TransferQueue does not automatically recreate the runtime.
• Storage-side PUT memory: PUT COMMIT copies received frames into ordinary memory before storing decoded values. Stored samples therefore do not keep the registered receive buffer pinned until CLEAR, and the buffer can return to the receive cache after the request finishes. This implements the storage-side copy-out part of the first memory-management suggestion. Client GET remains zero-copy, so long-lived GET results can still hold active receive leases. The alternative fixed, pre-registered object pool from the second suggestion has not been implemented.
• Abandoned PREPARE: Pending GET and PUT preparations expire. Expired GET state is released; an expired PUT's exposed receive buffer is quarantined because a late RDMA WRITE may still target it. Late COMMITs are rejected.
• Deferred response failure: If sending a deferred reply fails, the storage worker records a terminal failure. The proxy returns an explicit error for pending and later requests instead of forwarding them to a stopped worker. A StorageUnit-level test covers a deferred GET interleaved with another request and an injected send failure.

Validation of the pushed code tree: On each of two servers, the focused NIXL lifecycle, metrics, payload-transfer, and SimpleStorageUnit suites passed (190 passed, 2 warnings), and compileall passed. In a two-node Ray Actor setup, a second GET succeeded after the first completion reply was deliberately lost; 16 PUTs without CLEAR used one receive-MR registration and 15 reuses; and abandoned PUT/GET preparations expired as intended. UCX logs confirmed inter-node RMA over rc_verbs/hns_0:1 in all three runs.

Remaining scope: Client GET lease growth, partial-metadata reloads under varied buffer sizes, small-frame registration and lock contention, and large-payload copy cost still need workload-scale measurements. The quarantine budget stops new receives after its threshold; it is not a strict cap on a transfer already in flight.