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

docs/nixl_ucx_payload.md

@@ -76,6 +76,31 @@ or GID, add the corresponding variables to `ucx_env_vars`.
 If NIXL initialization or a transfer fails, TQ reports the error directly and does not fall back to ZMQ.
 If the transport is not restricted, or if `UCX_TLS` includes `tcp`, UCX may use TCP.
 
+### Receive Buffer Cache
+
+Set `payload_transfer.receive_buffer_cache_mb` to limit idle registered buffers **per runtime**,
+in MB (1024² bytes). The default is `256`; `0` disables caching. Larger payloads still transfer,
+but their buffers are released when no longer referenced. A storage-side PUT copies received
+frames into ordinary memory before storing decoded values, so stored values do not hold the
+registered receive buffer. Client GET views, in-flight transfers and fault-retained memory
+are outside the idle cache limit. A smaller cache saves idle memory at the cost of
+more registrations; account for all runtimes on a node when increasing it.
+
+Set `payload_transfer.quarantine_limit_mb` to stop preparing new NIXL receives after
+uncertain GETs have retained at least this many bytes per runtime. The default is `1024` MB.
+An already exposed receive buffer is never freed to satisfy the limit; it remains registered
+until the payload transfer closes. One in-flight receive can therefore take the retained total
+above the configured threshold. Increase the limit if normal payloads need more headroom.
+
+Set `payload_transfer.prepare_timeout_seconds` to expire abandoned PUT/GET PREPARE state.
+The default is twice the larger of the control-plane timeout and 180 seconds (normally
+400 seconds). The storage worker checks once per second. Expired GET frames are released
+and a late COMMIT is rejected. An expired PUT is removed from pending state, but its
+receive buffer must be quarantined until runtime teardown because the peer may still
+hold a READY token and perform a late WRITE. Quarantined bytes count against
+`quarantine_limit_mb`; do not deregister an exposed PUT buffer solely on elapsed time.
+If abandoned PUTs exhaust that limit, stop or restart the owning worker/job.
+
 ### Common UCX Configuration
 
 | Variable | Purpose | Reference value |
@@ -117,6 +142,27 @@ validation only confirm that the NIXL-UCX path is usable; to confirm RDMA, also
 
 ## Common Issues
 
+### A NIXL Session Has Failed
+
+An uncertain WRITE retains its buffers until teardown. The owning NIXL payload transfer enters
+a terminal state: the failed request raises, and all later PUT/GET requests through that transfer
+fail immediately, including requests to other peers. A storage unit sends a terminal error to
+the requesting client so its payload transfer also stops. Stop or restart the owning worker/job;
+TQ does not automatically rebuild a native agent after an uncertain WRITE. Confirmed safe
+pre-WRITE failures and local decode errors after completion do not poison the session. A
+completed WRITE with failed sender cleanup is terminal as well.
+The manager propagates `PayloadTransferFatalError` without retrying or wrapping it, so the job
+can distinguish this terminal failure from an ordinary request error.
+
+PUT cancellation waits for its running send. GET cancellation after COMMIT waits within the
+remaining timeout: a matching completion or pre-WRITE error releases the receiver; an unknown
+result quarantines only that receive buffer. A later GET uses a new transfer ID and buffer.
+Once the quarantine limit is reached, new receives fail explicitly until the payload transfer
+is recreated. Cancellation still propagates to the caller.
+
+With metrics enabled, `tq_storage_payload_*` reports cache, leased and quarantined bytes,
+queue wait, sender resources, and `terminal_send_failure` (0 or 1).
+
 ### RDMA Devices Are Ready, but NIXL-UCX Fails to Start
 
 If `ibv_devinfo` shows RDMA devices and active ports but NIXL initialization fails, the log typically contains:

tests/test_metrics.py

@@ -457,6 +457,36 @@ def mock_query(su_info, su_id):
         assert exporter.storage_capacity.labels(storage_unit_id="SU_002")._value.get() == 500
         assert call_count == 2
 
+    def test_payload_metrics_collected_and_removed_when_backend_no_longer_reports_them(self):
+        exporter = TQMetricsExporter()
+        exporter._storage_unit_infos = {"storage": MagicMock()}
+        snapshot = {
+            "storage_unit_id": "storage",
+            "payload_transfer": {"idle_receive_bytes": 128, "queue_wait_seconds": 2.5, "terminal_send_failure": 1},
+        }
+        exporter._query_storage_unit = MagicMock(return_value=snapshot)
+
+        exporter.collect_storage_metrics()
+        assert (
+            exporter.registry.get_sample_value("tq_storage_payload_idle_receive_bytes", {"storage_unit_id": "storage"})
+            == 128
+        )
+        assert (
+            exporter.registry.get_sample_value("tq_storage_payload_queue_wait_seconds", {"storage_unit_id": "storage"})
+            == 2.5
+        )
+        assert (
+            exporter.registry.get_sample_value("tq_storage_payload_terminal_send_failure", {"storage_unit_id": "storage"})
+            == 1
+        )
+
+        snapshot.pop("payload_transfer")
+        exporter.collect_storage_metrics()
+        assert (
+            exporter.registry.get_sample_value("tq_storage_payload_idle_receive_bytes", {"storage_unit_id": "storage"})
+            is None
+        )
+
     def test_storage_metrics_skips_capacity_when_none(self):
         """When capacity is None (unlimited storage), capacity/utilization
         gauges must not be populated, but other metrics still are.
@@ -567,11 +597,25 @@ def worker(self):
 
     @staticmethod
     def dispatch(worker, operation, body):
+        from queue import SimpleQueue
+
+        from transfer_queue.storage.payload_transfer import DeferredResponse
+        from transfer_queue.storage.simple_storage import _drain_deferred_responses
+
         unit, monitor = worker
         request = ZMQMessage.create(request_type=operation, sender_id="manager", body=body)
         socket = MagicMock()
         socket.recv_multipart.return_value = [b"manager"] + request.serialize()
-        unit._process_one_worker_request(socket, monitor)
+        measurements = {}
+        identity, response = unit._process_one_worker_request(socket, monitor, measurements)
+        if isinstance(response, DeferredResponse):
+            assert response.future.done()
+            completions = SimpleQueue()
+            completions.put((bytes(identity), response.future))
+            _drain_deferred_responses(completions, socket, unit.storage_unit_id, measurements)
+        else:
+            socket.send_multipart([identity] + response, copy=False)
+        assert measurements == {}
         sent = socket.send_multipart.call_args.args[0]
         assert sent[0] == b"manager"
         return ZMQMessage.deserialize(sent[1:])
@@ -609,22 +653,20 @@ def test_zmq_metrics_and_checkpoints_keep_their_arrival_labels(self, worker, tmp
     @pytest.mark.parametrize("finish", ["COMMIT", "CANCEL"])
     def test_nixl_counts_prepare_once_and_skips_continuations(self, worker, operation, finish):
         from transfer_queue.storage.payload_transfer.nixl import NixlPayloadTransfer, PayloadDescriptor
-        from transfer_queue.utils.serial_utils import calc_packed_size, encode, pack_into
+        from transfer_queue.utils.serial_utils import encode
 
         unit, monitor = worker
         unit.storage_data.put_data({"value": [42]}, [0])
         transfer = object.__new__(NixlPayloadTransfer)
         transfer._pending_puts = {}
         transfer._pending_gets = {}
         transfer._runtime = MagicMock()
+        transfer._runtime.terminal_error = None
         unit.payload_transfer = transfer
         frames = encode({"value": [42]})
-        payload = bytearray(calc_packed_size(frames))
-        pack_into(payload, frames)
-        descriptor = PayloadDescriptor("payload", len(payload), tuple(memoryview(frame).nbytes for frame in frames))
-        received = Future()
-        received.set_result(memoryview(payload))
-        transfer._runtime.receive.return_value = received
+        frame_sizes = tuple(memoryview(frame).nbytes for frame in frames)
+        descriptor = PayloadDescriptor("payload", sum(frame_sizes), frame_sizes)
+        transfer._runtime.receive.return_value = tuple(memoryview(frame) for frame in frames)
         sent = Future()
         sent.set_result(None)
         transfer._runtime.send.return_value = sent