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

[0oshowero0]

[AI Review]

Reviewed 73fcaa1f17f969b4694749e75ed7f03a486ad9ff against 6a863de8544a78e867105dc7a765876eb06589e4, including the October 13 review comments and October 16 follow-up fixes. The current diff is 14 files (+3182 / −363); the earlier concern about carrying the main-to-rdma synchronization is resolved.

The strategy boundary, frame-native scatter approach, storage-side PUT copy-out, and worker-owned deferred responses are reasonable. Several substantial historical findings are fixed. I would address the remaining lifecycle issues below before merging, and resolve the small-frame preparation bottleneck.

Findings to address before merge

1. [P1] Completed-WRITE cleanup failure can break native-agent-before-buffer teardown

Completed-WRITE cleanup, lines 647–665, interacting with close(), lines 430–456.

The failed-WRITE path deliberately clears its local handle, sources, and failed_sources before raising, so a retained Future/exception cannot keep the native handle alive. The DONE + release_xfer_handle() failure path omits that step:

except Exception as exc:
    self._retained_resources.append((handle, sources))
    cleanup_failed = True
...
if cleanup_failed:
    raise NixlWriteCompletedCleanupFailed(self._terminal_error)

The Future stores the exception; its traceback retains the _send_scatter frame and non-None handle. During close(), clearing the runtime's retained handles and deleting the agent wrapper therefore does not necessarily destroy the underlying native agent. The subsequent clearing of pending/quarantined receive owners can free exposed memory first.

This ownership relationship is present in the official NIXL implementation: the handle stores its agent, and initialize_xfer passes the underlying self.agent.

A reproduction using the actual PR runtime and a fake agent modeling that ownership produced:

exception: NixlWriteCompletedCleanupFailed
traceback retains native handle: True
after close: agent wrapper destroyed, receiver freed
native agent still alive: True
after dropping Future/exception: handle destroyed, native agent destroyed

With another pending or quarantined receive, this violates the teardown invariant relied upon to prevent late WRITE access to freed/reused memory. The reproduction establishes the ownership-ordering defect; actual RDMA memory corruption was not tested.

Please apply the same traceback-local cleanup discipline to this branch, preferably through one shared retention/raise exit. Add the missing combination test: DONE → handle cleanup failure → keep the Future/exception alive → close with an exposed receive → native agent must be destroyed before receive owners. The existing cleanup-failure, failed-Future-retention, and teardown-order tests cover these pieces separately.

2. [P2] Worker and proxy retain the last ZMQ GET response while idle

Worker, lines 1075–1089 and proxy, lines 996–1006.

Returning serialized frames fixes double serialization, but response now remains a local in the long-lived worker loop after send_multipart(). The proxy's newly introduced response = back.recv_multipart(...) does the same. Neither drops those frames before polling again.

Using actual SimpleStorage worker/proxy threads and local ZMQ sockets, an 8 MiB GET left these response-frame references after the client received and released its result:

Location	Frame bytes retained while idle	After the next CLEAR response
Worker	8,388,751	206
Proxy	8,388,777	232

The two references may share backing memory; these numbers do not establish two independent 8 MiB allocations. They do show that the final batch's buffers survive until another response overwrites the locals.

This affects the default ZMQ backend, particularly serialization copies, contiguous copies of non-contiguous tensors, and temporary SSD read results. Please drop the loop-local payload references after sending, or use a short-lived send scope. Add an idle-after-GET ownership test through the actual worker/proxy, in addition to counting serialization calls.

3. [P2] GET cancellation discards a confirmed remote terminal failure

Cancellation cleanup, lines 383–390, compared with normal response validation in _expect(), lines 710–723.

If a GET is cancelled after COMMIT, cleanup waits for the completion response and deserializes it, but never calls _expect() or otherwise records terminal_failure=True.

A remote WRITE can fail terminally during this window. The client correctly quarantines its receive buffer and rethrows CancelledError, but leaves _terminal_remote_error unset. Subsequent requests to other storage units pass _ensure_operational(), despite the documented runtime-wide failure contract.

Reproduced with the actual protocol and the repository's FakeAgent:

remote terminal=True
client terminal=None
next-operation guard accepts

Please share completion validation/state updates between normal and cancellation paths. Preserve cancellation as the outward exception while retaining the confirmed terminal state. Add a cancelled-GET test whose completion contains terminal_failure=True.

4. [P2, still outstanding] Frame preparation remains quadratic, with registration under the global lock

_acquire_frame_source(), lines 537–550.

Each writable frame scans leased receive buffers and all previously registered sources. For n disjoint ordinary source frames, the source-overlap scan performs n(n−1)/2 comparisons, followed by one registration per frame.

Calling the real source-acquisition method with FakeAgent and disjoint 8-byte frames:

Frames	Overlap comparisons	Local preparation time
2,048	2,096,128	0.169 s
4,096	8,386,560	0.885 s
8,192	33,550,336	2.335 s

These measurements exclude real MR registration cost and are not RDMA throughput measurements. Large frame counts are reachable: nested tensors are unbound, and storage GET returns per-sample tensor lists.

The lock is released between frames, but each scan and native registration holds the same runtime lock. Per-peer executors therefore do not remove preparation-stage contention. Please coalesce small frames and/or avoid the quadratic overlap scan and excessive registrations. This remains a material issue for a performance-oriented PR.

Historical review follow-up

I checked the code and tests against the original review, memory-management follow-up, and author's fix summary.

Earlier concern	Current status
Default ZMQ GET serialized twice	Fixed; serialized frames are reused, and a serialization-count test was added. The idle response-retention issue above is a separate regression.
An uncertain GET permanently disables that target	Fixed; only its receive buffer is quarantined, and a later independent GET uses a fresh transfer ID/buffer. The lost-completion/second-GET test covers this.
Unsafe native WRITE leaves a partially working runtime with no recovery owner	Normal paths now make the whole runtime terminal and propagate PayloadTransferFatalError unchanged. Automatic recreation is intentionally absent; the owning worker/job must stop or restart. Cancellation propagation still has finding 3.
Stored PUT data keeps the receive MR registered until CLEAR	Fixed on the storage side: each frame is copied to ordinary memory before decode/store. The repeated-PUT-without-CLEAR test checks MR reuse and data preservation.
Abandoned PREPARE state never expires	Implemented. Expired GET frames are released; exposed PUT receivers are quarantined, which is the safe choice because elapsed time alone cannot rule out a late WRITE.
Deferred reply failure silently stops the worker	Fixed for that path. Worker failure is recorded; the proxy answers pending and later requests explicitly. A real worker/proxy test covers an interleaved request and injected deferred-send failure.
remove/add remote agent whenever metadata changes	Still present. PUT copy-out improves MR reuse and reduces its frequency, but varying-size/concurrent/cache-miss workloads still need native validation. The code alone does not establish a QP leak or prove connection reuse.
Quadratic overlap checks / lock contention	Still present; independently reproduced above.
Main-to-rdma sync bundled into this PR	Resolved by the current base.
State-machine, metrics, and interface simplification	Mostly outstanding; follow-up fixes are separate commits, but the core feature still combines several substantial mechanisms.

Architecture, API, and maintainability

Keep the existing strategy and socket-ownership boundary. handle_request() -> ZMQMessage | DeferredResponse | None is a useful seam: background senders finish transfers while the worker remains the sole owner of its response socket. Storage-side PUT copy-out also gives communication buffers and stored values appropriately different lifetimes.

Describe GET zero-copy at the correct API boundary. The transport can return MR-backed views, but public get_data() scatters and repacks results. For ordinary pure-tensor fields, _pack_field_values() uses torch.stack or torch.nested.as_nested_tensor, creating new backing storage. A local check confirmed the packed batch does not share its input storage.

Thus the optimization can remove an intermediate transport copy; the standard pure-tensor public GET still copies during batch assembly. Long-lived receive leases remain relevant to lower-level calls and containers that preserve views. Also, after PUT copy-out, ordinary stored GET sources no longer originate in the PUT receive MR. Please benchmark through the public API to establish the actual benefit of leases and receive-MR-as-send-source reuse.

Reduce duplicated state handling. Findings 1 and 3 are examples of the same outcome being interpreted differently in separate cleanup paths. Make the pre-WRITE, posted-but-unconfirmed, and confirmed-complete phases explicit, and share their ownership/completion rules.

Concrete deletion opportunities remain:

• The per-peer _failed_peers checks are now redundant: _fail_sender() also sets _terminal_error, and preceding _ensure_open() calls already reject every terminal runtime. Keep diagnostic counts only if useful.
• Metric names and bookkeeping are spread across runtime diagnostics, strategy diagnostics, and an exporter whitelist. Centralize the exported schema and reconsider internal-only counters.
• GET COMMIT's receiver_endpoint is reconstructed from the token and checked against it again; it carries no independent information.
• initialize_packed_frame_table() has no production callers after this change.
• Update the main description to distinguish storage PUT copy-out, transport GET views, and public GET reassembly. The quarantine setting is a stop-new-receives threshold, not a strict cap on already exposed/in-flight memory.

A lower-priority edge case: an expired GET COMMIT currently reports write_not_started=False, so its client receiver is quarantined even when no WRITE occurred. This was reproduced, but conservative retention is safe and the default prepare TTL is long. Any optimization must distinguish a known expired preparation from an arbitrary unknown or possibly committed ID.

Validation

python -m compileall -q transfer_queue tutorial tests passed.

The following focused suites passed: test_nixl_payload_lifecycle, test_payload_transfer, test_metrics, test_storage_request_retry, test_async_simple_storage_manager, test_zmq_socket_pool, and test_simple_storage_unit.

263 passed, 2 warnings in 80.00 s.

Additional checks reproduced findings 1–3, measured the small-frame preparation path, and verified tensor reassembly copies. Native behavior used FakeAgent or an ownership model checked against official NIXL source. The ZMQ retention check used real worker/proxy threads and sockets. Real NIXL/UCX/RDMA hardware tests were not rerun locally; the author's two-node results remain separate supporting evidence.

[0oshowero0]

[AI Review] Follow-up: [P2] Receive deregistration failure permanently disables receives without a terminal error

Additional finding on 73fcaa1f17f969b4694749e75ed7f03a486ad9ff, supplementing the latest review.

When receive-memory deregistration fails, _deregister_receive_buffer() correctly retains the buffer in quarantine, but only sets _receive_cleanup_failed = True. _acquire_receive_buffer() then permanently rejects subsequent non-empty receives. Meanwhile, terminal_error remains None, so the request-entry terminal guards do not recognize that the receiver requires recreation.

This affects client-side GET or storage-side PUT. It can arise when returning an uncached receive buffer, evicting an idle buffer, or releasing the last receive lease. A worker can remain partially usable while every new receive fails as an ordinary request error, defeating the new PayloadTransferFatalError contract.

Reproduction

The following body can be added to tests/test_nixl_payload_lifecycle.py using its existing runtime / FakeAgent fixture and imports. It records the current incorrect behavior:

@pytest.mark.parametrize("runtime", [0], indirect=True)
def test_receive_cleanup_terminal_gap(runtime):
    instance, agent = runtime
    instance.prepare_receive(PayloadDescriptor("first", 16, (16,)))
    agent.deregister_error = True
    instance.cancel_receive("first")

    assert instance.diagnostics["quarantined_bytes"] == 16
    assert instance._receive_cleanup_failed is True
    assert instance.terminal_error is None

    with pytest.raises(NixlError) as raised:
        instance.prepare_receive(PayloadDescriptor("next", 16, (16,)))
    assert not isinstance(raised.value, PayloadTransferFatalError)

With the same injected failure, I also exercised the actual strategy's handle_request() and the manager's _get_from_single_storage_unit() method, bypassing socket leasing and supplying a protocol stub. The observed results were:

runtime: receive_cleanup_failed=True, terminal_error=None
next receive: NixlError, is_fatal=False
storage PUT_DATA_PREPARE: PUT_GET_ERROR, terminal_failure absent
manager GET: RuntimeError, is_fatal=False

The storage-side loss occurs because handle_request() / _handle_put_prepare() only recognizes runtime.terminal_error, and the ordinary error response carries no terminal flag. On the client side, the manager consequently takes its generic error-wrapping path.

Suggested fix and regression coverage

• Keep the failed registration and its backing buffer quarantined until native-agent teardown.
• Record this permanent receiver failure in the same terminal state checked by _ensure_open() and _ensure_operational(). Cleanup can run from a lease finalizer, so record the state there and surface the fatal error at request boundaries.
• Preserve the fatal classification through the generic prepare_receive() exception wrapper and the storage protocol's terminal_failure=True response. Cover a failure first discovered during the current PREPARE's cache eviction, after its initial terminal check has already passed.
• Extend the existing test_receive_cleanup_failure_retains_memory_and_stops_new_allocations to assert the terminal reason and fatal exception type, then check storage error-response propagation and manager preservation of PayloadTransferFatalError.

Validation here uses the PR's real Python lifecycle/protocol code with fault-injected FakeAgent behavior; native NIXL/UCX deregistration failure was not exercised on RDMA hardware.