From a4bbdd2e4fb97aec82d58591ba69494985377418 Mon Sep 17 00:00:00 2001 From: weiziyang Date: Tue, 8 Sep 2026 13:17:45 +0800 Subject: [PATCH] [optim] Lease pre-registered receive buffers on the MooncakeStore read path Tensor reads allocated a receive buffer and called register_buffer / unregister_buffer around every transfer. Registration is a kernel operation (page pinning + MR setup) whose throughput is an order of magnitude below the RDMA transfer it enables, so it dominated read time and left the RDMA backend slower than the TCP one. Lease the receive buffer from the local buffer that store.setup() already registered, then copy into the caller's tensors so the returned tensors keep owning their memory. split_by_bytes() bounds each lease by one reader thread's share of the pool, so a read no longer needs local_buffer_size to exceed the batch size. Builds without mooncake.store.BufferPool, and requests the pool cannot serve, still register their own receive buffers. Refs: https://github.com/Ascend/TransferQueue/issues/169 Signed-off-by: weiziyang --- tests/test_mooncake_client_lease.py | 206 ++++++++++++++++++ .../storage/clients/mooncake_client.py | 129 ++++++++++- 2 files changed, 333 insertions(+), 2 deletions(-) create mode 100644 tests/test_mooncake_client_lease.py diff --git a/tests/test_mooncake_client_lease.py b/tests/test_mooncake_client_lease.py new file mode 100644 index 00000000..4d27c344 --- /dev/null +++ b/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 == [] diff --git a/transfer_queue/storage/clients/mooncake_client.py b/transfer_queue/storage/clients/mooncake_client.py index d6914902..0fd85b2b 100644 --- a/transfer_queue/storage/clients/mooncake_client.py +++ b/transfer_queue/storage/clients/mooncake_client.py @@ -40,6 +40,13 @@ 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 @@ -47,6 +54,48 @@ 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