[feat] introduce payload transfer abstraction and NIXL-UCX backend

README.md

@@ -211,6 +211,9 @@ pip install TransferQueue
    pip install dist/*.whl
    ```
 
+For the optional SimpleStorage NIXL-UCX Host payload path, see
+[the NIXL-UCX payload guide](docs/nixl_ucx_payload.md).
+
 <h2 id="performance">📊 Performance</h2>
 
 ### Simple Case: Regular Tensor
@@ -345,4 +348,4 @@ Please kindly cite our paper if you find this repo is useful:
   journal={arXiv preprint arXiv:2507.01663},
   year={2025}
 }
-```
+```

docs/nixl_ucx_payload.md

@@ -0,0 +1,188 @@
+# SimpleStorage NIXL-UCX Host Payload Transfer
+
+After setting `payload_transfer` to `nixl-ucx`, all non-empty payloads are transferred through NIXL-UCX;
+ZMQ handles control messages only. This document describes the configuration and usage.
+
+## 1. Check RDMA Devices
+
+Run the following on every node running TQ or a Ray worker:
+
+```bash
+ls /sys/class/infiniband
+rdma link show
+ibv_devinfo
+```
+
+`ls` should list RDMA devices, and the ports shown by `rdma link show` should be `ACTIVE`.
+`ibv_devinfo` does not show the provider dynamic library name. If no device is present or a port is
+not active, check the driver, `rdma-core`, provider, and container device mappings first.
+
+## 2. Install NIXL and TQ
+
+Install the NIXL wheel:
+
+```bash
+python -m pip install nixl
+```
+
+The NIXL wheel includes the UCX runtime, but the system must still provide `rdma-core`,
+`libibverbs`, and the provider for the network adapter.
+
+Install TQ from the source directory:
+
+```bash
+python -m pip install -e .
+```
+
+## 3. Check the NIXL-UCX Backend
+
+Run:
+
+```bash
+python - <<'PY'
+from nixl import nixl_agent, nixl_agent_config
+
+config = nixl_agent_config(
+    enable_prog_thread=True,
+    enable_listen_thread=True,
+    listen_port=0,
+    backends=["UCX"],
+)
+agent = nixl_agent("tq-nixl-check", config)
+assert "UCX" in agent.backends
+print("NIXL UCX backend is available")
+PY
+```
+
+The command should output `NIXL UCX backend is available`. If it fails, see the common issues at the end.
+
+## 4. Enable SimpleStorage NIXL-UCX Transfer
+
+Enable NIXL-UCX in the TQ configuration:
+
+```yaml
+backend:
+  storage_backend: SimpleStorage
+  SimpleStorage:
+    payload_transfer:
+      backend: nixl-ucx
+      ucx_env_vars: {}
+```
+
+`ucx_env_vars: {}` means that TQ does not set additional UCX environment variables. TQ and Ray workers
+continue to use the `UCX_*` variables inherited when they were started. To specify a transport, device,
+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.
+
+### Common UCX Configuration
+
+| Variable | Purpose | Reference value |
+| --- | --- | --- |
+| `UCX_TLS` | Restrict the transports available to UCX | `<available rc_* transport>,tcp,sm,self` |
+| `UCX_NET_DEVICES` | Specify the RDMA device and port | `<rdma_device>:<port>` |
+| `UCX_IB_GID_INDEX` | Specify the RoCE GID index | `<gid_index>` |
+| `UCX_MODULE_DIR` | Specify the UCX transport module directory in the NIXL wheel | `<ucx_module_dir>` |
+
+Restart TQ/Ray after making changes. Set the device name and GID index for each node.
+
+### Memory Registration
+
+Before NIXL registers memory, check the system limit in the current shell:
+
+```bash
+ulimit -l
+```
+
+If the value is too small, set it to `unlimited` in the shell that starts TQ/Ray:
+
+```bash
+ulimit -l unlimited
+```
+
+This setting applies only to the current shell and its child processes.
+
+## 5. Verify SimpleStorage NIXL-UCX Transfer
+
+After enabling it, the StorageUnit startup log will contain:
+
+```text
+SimpleStorage payload transfer selected: nixl-ucx device=ucx-auto gid_index=ucx-auto tls=ucx-auto
+```
+
+After a cross-node PUT/GET completes, the GET content should match the PUT content. The log and data
+validation only confirm that the NIXL-UCX path is usable; to confirm RDMA, also check the payload lane.
+`rc_*` indicates RDMA, while a TCP lane indicates that TCP is being used.
+
+## Common Issues
+
+### 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:
+
+```text
+no userspace device-specific driver found
+failed to open ... libuct_ib ...
+NIXL_ERR_BACKEND
+```
+
+First confirm that the provider for the network adapter is installed. If the provider is installed but the
+error persists, use a NIXL wheel compatible with the system `rdma-core/provider`. If no suitable wheel is
+available, follow the [official NIXL source build instructions](https://github.com/ai-dynamo/nixl#prerequisites-for-source-build-linux)
+to build UCX with multi-thread and verbs enabled:
+
+```bash
+python -m pip install meson ninja pybind11 tomlkit
+git clone https://github.com/openucx/ucx.git <ucx_source>
+cd <ucx_source>
+git checkout <nixl_supported_ucx_version>
+./autogen.sh
+./contrib/configure-release-mt \
+  --prefix=<ucx_install_prefix> \
+  --enable-shared \
+  --disable-static \
+  --with-verbs
+make -j"$(nproc)"
+make install
+```
+
+Then configure NIXL to use this UCX:
+
+```bash
+git clone https://github.com/ai-dynamo/nixl.git <nixl_source>
+cd <nixl_source>
+python -m pip install .
+meson setup build \
+  -Ducx_path=<ucx_install_prefix> \
+  -Dprefix=<nixl_install_prefix> \
+  -Dbuildtype=release
+ninja -C build
+ninja -C build install
+python -m pip install build/src/bindings/python/nixl-meta/nixl-*-py3-none-any.whl
+```
+
+### `libnixl.so` Cannot Be Found
+
+If `import nixl` reports `libnixl.so: cannot open shared object file`, the NIXL Python extension
+cannot find the NIXL shared library; UCX initialization has not started yet.
+
+First locate `libnixl.so` in the wheel:
+
+```bash
+NIXL_SITE=$(python -c 'import site; print(site.getsitepackages()[0])')
+find "${NIXL_SITE}" -name libnixl.so
+```
+
+If it is found, add its containing directory to `LD_LIBRARY_PATH` in the same shell that starts TQ/Ray:
+
+```bash
+NIXL_LIB_DIR=/path/to/directory/containing/libnixl.so
+export LD_LIBRARY_PATH="${NIXL_LIB_DIR}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
+```
+
+If it is not found, reinstall the NIXL wheel corresponding to the reported error. For example:
+
+```bash
+python -m pip install --no-cache-dir --force-reinstall --no-deps nixl-cu12
+```

tests/test_payload_transfer.py

@@ -0,0 +1,146 @@
+# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Payload transfer contract and NIXL-UCX configuration tests."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+from omegaconf import OmegaConf
+
+from transfer_queue.storage.payload_transfer import (
+    PayloadTransferError,
+    create_payload_transfer,
+    parse_payload_transfer_config,
+)
+from transfer_queue.storage.payload_transfer.nixl import PayloadDescriptor
+from transfer_queue.storage.payload_transfer.nixl_ucx_runtime import _configure_ucx_environment
+from transfer_queue.storage.payload_transfer.zmq import ZmqPayloadTransfer
+from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType
+
+
+def test_payload_descriptor_preserves_frame_layout():
+    descriptor = PayloadDescriptor("framed", 4 + 8 * 2 + 5, (2, 3))
+    descriptor.validate()
+    assert PayloadDescriptor.from_dict(descriptor.to_dict()) == descriptor
+
+    with pytest.raises(PayloadTransferError, match="packed payload length"):
+        PayloadDescriptor("framed", 5, (2, 3)).validate()
+
+
+def test_payload_descriptor_requires_frame_layout_and_rejects_negative_lengths():
+    with pytest.raises(KeyError, match="frame_sizes"):
+        PayloadDescriptor.from_dict({"transfer_id": "payload", "payload_bytes": 3})
+
+    with pytest.raises(PayloadTransferError, match="negative"):
+        PayloadDescriptor.from_dict({"transfer_id": "payload", "payload_bytes": -1, "frame_sizes": [1]})
+
+
+def test_yaml_ucx_settings_override_process_environment(monkeypatch):
+    monkeypatch.setenv("UCX_TLS", "sm")
+    monkeypatch.delenv("UCX_IB_GID_INDEX", raising=False)
+
+    configured = _configure_ucx_environment(
+        {
+            "UCX_TLS": "tcp",
+            "UCX_IB_GID_INDEX": 3,
+        }
+    )
+
+    assert configured == {"UCX_TLS": "tcp", "UCX_IB_GID_INDEX": "3"}
+    assert os.environ["UCX_TLS"] == "tcp"
+    assert os.environ["UCX_IB_GID_INDEX"] == "3"
+
+
+def test_empty_yaml_ucx_settings_preserve_process_environment(monkeypatch):
+    monkeypatch.setenv("UCX_NET_DEVICES", "custom_hca:2")
+
+    assert _configure_ucx_environment({}) == {}
+    assert os.environ["UCX_NET_DEVICES"] == "custom_hca:2"
+
+
+def test_payload_transfer_rejects_unsupported_backend():
+    with pytest.raises(ValueError, match="expected 'zmq' or 'nixl-ucx'"):
+        create_payload_transfer({"backend": "unsupported"})
+
+
+def test_factory_returns_zmq_payload_transfer():
+    transfer = create_payload_transfer({"backend": "zmq"})
+
+    assert isinstance(transfer, ZmqPayloadTransfer)
+    assert transfer.bootstrap_info() is None
+
+
+def test_nixl_factory_rejects_incomplete_peer_endpoints():
+    with pytest.raises(RuntimeError, match="payload transfer endpoints are missing"):
+        create_payload_transfer(
+            {"backend": "nixl-ucx"},
+            control_peer_infos={"storage": object()},
+        )
+
+
+def test_zmq_payload_transfer_handles_put_and_get_requests():
+    stored = {}
+    transfer = ZmqPayloadTransfer()
+
+    put_request = ZMQMessage.create(
+        request_type=ZMQRequestType.PUT_DATA,
+        sender_id="manager",
+        receiver_id="storage",
+        body={"global_indexes": [1], "data": {"value": [42]}},
+    )
+    put_response = transfer.handle_request(
+        put_request,
+        storage_id="storage",
+        load_data=lambda fields, indexes: {field: stored[field] for field in fields},
+        store_data=lambda indexes, data, parser: stored.update(data),
+    )
+
+    assert put_response.request_type == ZMQRequestType.PUT_DATA_RESPONSE
+    assert stored == {"value": [42]}
+
+    get_request = ZMQMessage.create(
+        request_type=ZMQRequestType.GET_DATA,
+        sender_id="manager",
+        receiver_id="storage",
+        body={"global_indexes": [1], "fields": ["value"]},
+    )
+    get_response = transfer.handle_request(
+        get_request,
+        storage_id="storage",
+        load_data=lambda fields, indexes: {field: stored[field] for field in fields},
+        store_data=lambda indexes, data, parser: stored.update(data),
+    )
+
+    assert get_response.request_type == ZMQRequestType.GET_DATA_RESPONSE
+    assert get_response.body["data"] == {"value": [42]}
+
+
+def test_payload_transfer_config_extracts_ucx_settings():
+    assert parse_payload_transfer_config(
+        {
+            "backend": "nixl-ucx",
+            "ucx_env_vars": {"UCX_TLS": "rc", "UCX_IB_GID_INDEX": 3},
+        }
+    ) == ("nixl-ucx", {"ucx_env_vars": {"UCX_TLS": "rc", "UCX_IB_GID_INDEX": 3}})
+
+
+def test_public_config_defaults_to_zmq_payload_transfer():
+    config = OmegaConf.load(Path(__file__).parents[1] / "transfer_queue/config.yaml")
+
+    assert config.backend.SimpleStorage.payload_transfer.backend == "zmq"

tests/test_serial_utils_batch_on_cpu.py

@@ -22,6 +22,8 @@
 * ``batch_decode_from``
 """
 
+import struct
+
 import numpy as np
 import pytest
 import torch
@@ -42,6 +44,15 @@ def test_calc_packed_size_then_pack_unpack_roundtrip():
     assert [bytes(mv) for mv in recovered] == items
 
 
+def test_initialize_packed_frame_table_leaves_payload_for_direct_receive():
+    items = [b"hello", b"world!"]
+    buf = bytearray(serial_utils.calc_packed_size(items))
+    serial_utils.initialize_packed_frame_table(buf, [len(item) for item in items])
+    payload_start = serial_utils._PACK_HEADER_SIZE + len(items) * serial_utils._PACK_ENTRY_SIZE
+    buf[payload_start:] = b"helloworld!"
+    assert [bytes(mv) for mv in serial_utils.unpack_from(buf)] == items
+
+
 def test_pack_into_writes_only_within_its_slice():
     items = [b"alpha", b"beta", b"gamma"]
     sz = serial_utils.calc_packed_size(items)
@@ -64,6 +75,17 @@ def test_unpack_from_zero_item_buffer():
     assert serial_utils.unpack_from(buf) == []
 
 
+def test_unpack_from_rejects_invalid_frame_bounds():
+    items = [b"payload"]
+    buf = bytearray(serial_utils.calc_packed_size(items))
+    serial_utils.pack_into(buf, items)
+
+    # Corrupt the frame offset so it points into the frame table.
+    struct.pack_into("<I", buf, serial_utils._PACK_HEADER_SIZE, 4)
+    with pytest.raises(ValueError, match="outside the payload"):
+        serial_utils.unpack_from(buf)
+
+
 # ============================================================================
 # batch_encode_into + batch_decode_from (high-level batch layer)
 # ============================================================================