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