[optim] Lease pre-registered receive buffers on the MooncakeStore read path

tests/test_mooncake_client_lease.py

@@ -0,0 +1,206 @@
+# 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.
+
+"""Receive-buffer leasing on the MooncakeStore tensor read path.
+
+The store and the lease pool are faked, so these tests need neither mooncake nor RDMA.
+"""
+
+import ctypes
+import sys
+
+import pytest
+import torch
+
+from transfer_queue.storage.clients import mooncake_client as mcc
+
+DTYPES = [torch.float32, torch.int64, torch.float32, torch.int16]
+SHAPES = [(4, 3), (5,), (), (2, 8)]
+KEYS = ["k0", "k1", "k2", "k3"]
+
+
+def expected_tensors():
+    """Deterministic payloads, one per key in KEYS."""
+    out = []
+    for seed, (dtype, shape) in enumerate(zip(DTYPES, SHAPES, strict=True)):
+        numel = torch.empty(shape).numel()
+        values = torch.arange(seed * 100, seed * 100 + numel)
+        out.append(values.to(dtype).reshape(shape))
+    return out
+
+
+class FakeStore:
+    """Writes the stored payload into whatever pointer batch_get_into is given."""
+
+    def __init__(self):
+        self.objects = {
+            key: bytes(t.contiguous().numpy().tobytes()) for key, t in zip(KEYS, expected_tensors(), strict=True)
+        }
+        self.registered: list[tuple[int, int]] = []
+        self.unregistered: list[int] = []
+
+    def setup(self, *args):
+        return 0
+
+    def register_buffer(self, ptr, size):
+        self.registered.append((ptr, size))
+        return 0
+
+    def unregister_buffer(self, ptr):
+        self.unregistered.append(ptr)
+        return 0
+
+    def batch_get_into(self, keys, ptrs, sizes):
+        for key, ptr, size in zip(keys, ptrs, sizes, strict=True):
+            ctypes.memmove(ptr, self.objects[key], size)
+        return list(sizes)
+
+
+class FakeLease:
+    def __init__(self, pool, nbytes):
+        self._pool = pool
+        self._memory = torch.empty(nbytes, dtype=torch.uint8)
+        self.ptr = self._memory.data_ptr()
+        # A numpy array (not its .data memoryview): torch.frombuffer keeps a reference
+        # to it, so a staged view left alive at release() shows up as an extra refcount.
+        self.buffer = self._memory.numpy()
+        self._baseline_refs = sys.getrefcount(self.buffer)
+
+    def release(self):
+        # Mirror mooncake: the pool refuses to return a lease while a view of its
+        # buffer is still alive (a live torch.frombuffer tensor holds a reference).
+        if sys.getrefcount(self.buffer) > self._baseline_refs:
+            raise RuntimeError("cannot release buffer while exported views exist")
+        self._pool.released += 1
+
+
+class FakePool:
+    """Serves leases up to ``capacity`` bytes; larger requests cannot be served."""
+
+    def __init__(self, capacity=1 << 30):
+        self.capacity = capacity
+        self.acquired: list[int] = []
+        self.released = 0
+
+    def acquire(self, nbytes, block=True):
+        if nbytes > self.capacity:
+            return None
+        self.acquired.append(nbytes)
+        return FakeLease(self, nbytes)
+
+
+@pytest.fixture
+def store(monkeypatch):
+    fake = FakeStore()
+    monkeypatch.setattr(mcc, "MOONCAKE_STORE_IMPORTED", True)
+    # raising=False: these symbols are absent unless mooncake is installed.
+    monkeypatch.setattr(mcc, "MooncakeDistributedStore", lambda: fake, raising=False)
+    monkeypatch.setattr(mcc, "ReplicateConfig", type("ReplicateConfig", (), {}), raising=False)
+    return fake
+
+
+def make_client(local_buffer_size=1 << 30):
+    return mcc.MooncakeStoreClient(
+        {
+            "local_hostname": "127.0.0.1",
+            "metadata_server": "127.0.0.1:8080",
+            "master_server_address": "127.0.0.1:8081",
+            "local_buffer_size": local_buffer_size,
+        }
+    )
+
+
+def read_all(client):
+    tensors, indexes = client._get_tensors_thread_worker(KEYS, SHAPES, DTYPES, list(range(len(KEYS))))
+    assert indexes == list(range(len(KEYS)))
+    return tensors
+
+
+def assert_payloads(tensors):
+    for got, want in zip(tensors, expected_tensors(), strict=True):
+        assert got.dtype == want.dtype
+        assert got.shape == want.shape
+        assert torch.equal(got, want)
+
+
+def install_pool(monkeypatch, capacity=1 << 30):
+    """Make the client believe mooncake provides a lease pool, and hand it a fake one."""
+    fake = FakePool(capacity)
+    monkeypatch.setattr(mcc, "MOONCAKE_BUFFER_POOL_IMPORTED", True)
+    monkeypatch.setattr(mcc, "BufferPool", lambda _store, max_bytes: fake, raising=False)
+    return fake
+
+
+def test_reads_land_in_leased_buffer(store, monkeypatch):
+    pool = install_pool(monkeypatch)
+
+    tensors = read_all(make_client())
+
+    assert_payloads(tensors)
+    # The whole point: no registration on the data path, and no lease left behind.
+    assert store.registered == []
+    assert len(pool.acquired) == 1 and pool.released == 1
+
+
+def test_batch_larger_than_lease_share_is_read_in_rounds(store, monkeypatch):
+    pool = install_pool(monkeypatch)
+
+    # Small local buffer: each thread's share holds only part of the batch.
+    client = make_client(local_buffer_size=256 * mcc.MAX_BATCH_WORKER_THREADS)
+    tensors = read_all(client)
+
+    assert_payloads(tensors)
+    assert len(pool.acquired) > 1
+    assert all(nbytes <= client._lease_bytes for nbytes in pool.acquired)
+    assert pool.released == len(pool.acquired)
+    assert store.registered == []
+
+
+def test_falls_back_to_own_buffers_when_pool_cannot_serve(store, monkeypatch):
+    pool = install_pool(monkeypatch, capacity=0)
+
+    tensors = read_all(make_client())
+
+    assert_payloads(tensors)
+    assert pool.acquired == []
+    assert store.registered and len(store.unregistered) == len(store.registered)
+
+
+def test_registers_receive_regions_without_lease_support(store, monkeypatch):
+    monkeypatch.setattr(mcc, "MOONCAKE_BUFFER_POOL_IMPORTED", False)
+
+    tensors = read_all(make_client())
+
+    assert_payloads(tensors)
+    assert store.registered and len(store.unregistered) == len(store.registered)
+
+
+def test_uniform_group_copied_in_one_strided_pass(store, monkeypatch):
+    # Many identical small tensors (the fragmented-read case) take the single strided
+    # copy-out path instead of a per-tensor loop; the payloads must still round-trip.
+    pool = install_pool(monkeypatch)
+    n = 8
+    dtypes = [torch.float32] * n
+    shapes = [(16,)] * n
+    keys = [f"u{i}" for i in range(n)]
+    payloads = [torch.arange(i, i + 16, dtype=torch.float32) for i in range(n)]
+    store.objects = {k: bytes(t.numpy().tobytes()) for k, t in zip(keys, payloads, strict=True)}
+
+    tensors, indexes = make_client()._get_tensors_thread_worker(keys, shapes, dtypes, list(range(n)))
+
+    assert indexes == list(range(n))
+    for got, want in zip(tensors, payloads, strict=True):
+        assert torch.equal(got, want)
+    assert pool.acquired and pool.released == len(pool.acquired)
+    assert store.registered == []

transfer_queue/storage/clients/mooncake_client.py

@@ -40,13 +40,62 @@
 except ImportError:
     MOONCAKE_STORE_IMPORTED = False
 
+MOONCAKE_BUFFER_POOL_IMPORTED: bool = True
+try:
+    from mooncake.store import BufferPool
+except ImportError:
+    # Older mooncake builds have no lease API; those fall back to registering per transfer.
+    MOONCAKE_BUFFER_POOL_IMPORTED = False
+
 BATCH_SIZE_LIMIT: int = 400
 MAX_BATCH_WORKER_THREADS = 4
 MAX_SERIAL_WORKER_THREADS = 4
 MAX_RETRIES = 3
 RETRY_DELAY_SECONDS = 1.0
 
 
+def _copy_lease_into_tensors(lease, targets: list[Tensor], offsets: list[int]) -> None:
+    """Copy each staged region out of the lease buffer into the caller's tensors.
+
+    A uniform, tightly-packed group (same dtype+shape, laid out contiguously in the
+    target region) is copied in a single strided pass. This is the many-small-key
+    case, where a per-tensor Python loop otherwise dominates the read; ragged groups
+    fall back to a per-tensor copy. Every frombuffer view is dropped before returning:
+    the pool refuses to release a lease while an exported view of its buffer is alive.
+    """
+    t0 = targets[0]
+    numel = t0.numel()
+    uniform = len(targets) > 1
+    if uniform:
+        base_off = t0.storage_offset()
+        base_ptr = t0.untyped_storage().data_ptr()
+        for j, t in enumerate(targets):
+            if (
+                t.dtype != t0.dtype
+                or t.shape != t0.shape
+                or not t.is_contiguous()
+                or t.storage_offset() != base_off + j * numel
+                or t.untyped_storage().data_ptr() != base_ptr
+            ):
+                uniform = False
+                break
+
+    if uniform:
+        k = len(targets)
+        stride_elems = (offsets[1] - offsets[0]) // t0.element_size()
+        src = torch.frombuffer(
+            lease.buffer, dtype=t0.dtype, count=(k - 1) * stride_elems + numel
+        ).as_strided((k, numel), (stride_elems, 1))
+        t0.as_strided((k, numel), (numel, 1), t0.storage_offset()).copy_(src)
+        del src
+        return
+
+    for target, off in zip(targets, offsets, strict=True):
+        staged = torch.frombuffer(lease.buffer, dtype=target.dtype, count=target.numel(), offset=off)
+        target.copy_(staged.view(target.shape))
+        del staged
+
+
 @StorageClientFactory.register("MooncakeStoreClient")
 class MooncakeStoreClient(StorageKVClient):
     """
@@ -134,6 +183,21 @@ def __init__(self, config: dict[str, Any]):
         if ret != 0:
             raise RuntimeError(f"Mooncake store setup failed with error code: {ret}")
 
+        # RDMA can only target registered (pinned) memory, and register_buffer is a kernel
+        # operation that costs far more than the transfer it enables. Lease receive buffers
+        # from the local buffer that setup() already registered instead of registering per
+        # transfer. See https://github.com/Ascend/TransferQueue/issues/169
+        # max_bytes=0: lease from that local buffer only, without an extra arena.
+        self._buffer_pool = BufferPool(self._store, max_bytes=0) if MOONCAKE_BUFFER_POOL_IMPORTED else None
+        if self._buffer_pool is None:
+            logger.warning(
+                "mooncake.store.BufferPool is unavailable, so every tensor read registers and "
+                "unregisters its own receive buffer. Upgrade mooncake-transfer-engine to lease "
+                "pre-registered buffers instead."
+            )
+        # One share per reader thread, so all of them can hold a lease at the same time.
+        self._lease_bytes = self.local_buffer_size // MAX_BATCH_WORKER_THREADS
+
     def put(self, keys: list[str], values: list[Any]) -> list[dict | None]:
         """Stores multiple key-value pairs to MooncakeStore.
 
@@ -408,13 +472,70 @@ def _get_tensors_thread_worker(
             batch_dtypes, batch_shapes
         )
 
+        if self._buffer_pool is None:
+            self._read_into_own_buffers(batch_keys, batch_buffer_ptrs, batch_nbytes, region_ptrs, region_sizes)
+            return batch_buffer_tensors, indexes
+
+        # split_by_bytes() keeps every lease request within one thread's share of the pool,
+        # so a batch larger than that share is read in several rounds instead of failing.
+        for group in split_by_bytes(batch_nbytes, self._lease_bytes):
+            self._read_group_via_lease(group, batch_keys, batch_buffer_ptrs, batch_nbytes, batch_buffer_tensors)
+
+        return batch_buffer_tensors, indexes
+
+    def _read_into_own_buffers(
+        self, keys: list[str], ptrs: list[int], nbytes: list[int], region_ptrs: list[int], region_sizes: list[int]
+    ) -> None:
+        """Register the receive regions for one transfer, read into them, then unregister."""
         self._register_all_buffers(region_ptrs, region_sizes)
         try:
-            self._batch_get_into_with_retry(batch_keys, batch_buffer_ptrs, batch_nbytes)
+            self._batch_get_into_with_retry(keys, ptrs, nbytes)
         finally:
             self._unregister_all_buffers(region_ptrs)
 
-        return batch_buffer_tensors, indexes
+    def _read_group_via_lease(
+        self,
+        group: list[int],
+        batch_keys: list[str],
+        batch_ptrs: list[int],
+        batch_nbytes: list[int],
+        batch_tensors: list[Tensor],
+    ) -> None:
+        """Read one group of keys into leased memory, then copy into the caller's tensors.
+
+        The copy is what lets the lease return to the pool immediately, keeping the
+        returned tensors owned by the caller exactly as the register-per-read path does.
+        """
+        keys = [batch_keys[i] for i in group]
+        nbytes = [batch_nbytes[i] for i in group]
+        offsets, total = _aligned_offsets(nbytes)
+
+        lease = self._acquire_lease(total)
+        if lease is None:
+            ptrs = [batch_ptrs[i] for i in group]
+            region_ptrs, region_sizes = merge_contiguous_memory(ptrs, nbytes)
+            self._read_into_own_buffers(keys, ptrs, nbytes, region_ptrs, region_sizes)
+            return
+
+        try:
+            self._batch_get_into_with_retry(keys, [lease.ptr + off for off in offsets], nbytes)
+            _copy_lease_into_tensors(lease, [batch_tensors[i] for i in group], offsets)
+        finally:
+            lease.release()
+
+    def _acquire_lease(self, nbytes: int):
+        """Lease ``nbytes`` of pre-registered memory, or None when the pool cannot serve it.
+
+        Never block: mooncake would otherwise wait for capacity that a request larger than
+        the local buffer never gets. Exhaustion surfaces as None in some builds and as an
+        exception in others, and both mean the caller should register its own memory.
+        """
+        assert self._buffer_pool is not None
+        try:
+            return self._buffer_pool.acquire(nbytes, block=False)
+        except Exception as e:
+            logger.warning(f"Leasing {nbytes} B of pre-registered memory failed ({e}); registering own buffer.")
+            return None
 
     def _get_tensors_gdr(
         self,
@@ -535,6 +656,10 @@ def clear(self, keys: list[str], custom_backend_meta: list[Any] | None = None) -
 
     def close(self):
         """Closes MooncakeStore."""
+        # Release the leased regions before the store they belong to goes away.
+        if self._buffer_pool is not None:
+            self._buffer_pool.close()
+            self._buffer_pool = None
         if self._gdr_staging is not None:
             self._gdr_staging.close(self._store)
             self._gdr_staging = None