From a4bbdd2e4fb97aec82d58591ba69494985377418 Mon Sep 17 00:00:00 2001 From: weiziyang Date: Tue, 8 Sep 2026 13:17:45 +0800 Subject: [PATCH 1/2] [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 From 7b5f93b70f62f64cf6c4152851c4eed3d871fbe1 Mon Sep 17 00:00:00 2001 From: weiziyang Date: Sat, 26 Sep 2026 21:17:32 +0800 Subject: [PATCH 2/2] [optim] Bound pooled receive buffers to the registered local buffer Review follow-up on the pooled MooncakeStore read path: - Pass max_bytes=local_buffer_size and max_regions=MAX_BATCH_WORKER_THREADS. On the default max_bytes=0 mooncake budgets twice the local buffer and serves the excess by registering fresh memory per lease, which is the cost this path removes; a value below the buffer is clamped back up, so the per-thread share is what reserves headroom. - Halve that share and send a group larger than it down the register-per-read path: split_by_bytes() gives an oversized tensor its own group without shrinking it, and the store stages its own reads from the same region. - Drop staged views on the copy-out error path, keep the original exception when returning a buffer fails, and refuse to close the store while a buffer is out. - Warn once instead of per group per read, and only where CPU RDMA reads pay for registration. - Name the helpers and the test module after the buffer pool instead of leases, and make the fake pool raise on exhaustion the way mooncake does. Signed-off-by: weiziyang --- tests/test_mooncake_buffer_pool.py | 458 ++++++++++++++++++ tests/test_mooncake_client_lease.py | 206 -------- .../storage/clients/mooncake_client.py | 110 +++-- 3 files changed, 528 insertions(+), 246 deletions(-) create mode 100644 tests/test_mooncake_buffer_pool.py delete mode 100644 tests/test_mooncake_client_lease.py diff --git a/tests/test_mooncake_buffer_pool.py b/tests/test_mooncake_buffer_pool.py new file mode 100644 index 00000000..d8d05a7b --- /dev/null +++ b/tests/test_mooncake_buffer_pool.py @@ -0,0 +1,458 @@ +# 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. + +"""Pooled receive buffers on the MooncakeStore tensor read path. + +The store and the buffer pool are faked, so these tests need neither mooncake nor RDMA. +The fake pool keeps the upstream properties this path depends on: buffers are carved out +of one fixed arena and reused, exhaustion raises, a buffer cannot be returned while an +exported view of it is alive, and the pool cannot close while a buffer is out on loan. +""" + +import ctypes +import logging +import threading +import weakref + +import pytest +import torch + +from transfer_queue.storage.clients import mooncake_client as mcc + +LOGGER_NAME = "transfer_queue.storage.clients.mooncake_client" +ALIGN = 256 + +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] = [] + self.closed = False + self.transfer_error: Exception | None = None + self.on_transfer = None + + 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): + if self.transfer_error is not None: + raise self.transfer_error + if self.on_transfer is not None: + self.on_transfer() + for key, ptr, size in zip(keys, ptrs, sizes, strict=True): + ctypes.memmove(ptr, self.objects[key], size) + return list(sizes) + + def close(self): + self.closed = True + + +class FakeBuffer: + """One region carved out of the pool's arena.""" + + def __init__(self, pool, offset, nbytes): + self._pool = pool + self._offset = offset + self.nbytes = nbytes + self.ptr = pool.arena.data_ptr() + offset + self._views: list[weakref.ref] = [] + self.returned = False + + @property + def buffer(self): + # Mooncake hands out a fresh view per access and counts it as an export. + view = self._pool.arena_bytes[self._offset : self._offset + max(self.nbytes, 1)] + self._views.append(weakref.ref(view)) + return view + + def has_live_views(self): + return any(ref() is not None for ref in self._views) + + def release(self): + if self.has_live_views(): + raise RuntimeError("cannot release buffer while exported views exist") + self._pool._give_back(self._offset) + self.returned = True + + +class FakePool: + """Sub-allocates one arena, the way mooncake carves up its registered local buffer.""" + + def __init__(self, capacity, max_regions=None): + self.capacity = capacity + self.max_regions = max_regions + self.arena = torch.empty(max(capacity, 1), dtype=torch.uint8) + self.arena_bytes = self.arena.numpy() + self.acquired: list[int] = [] + self.offsets: list[int] = [] + self.returned = 0 + self.active: dict[int, int] = {} + self.peak_bytes = 0 + self.peak_regions = 0 + self.closed = False + self._lock = threading.Lock() + + def acquire(self, nbytes, block=True): + with self._lock: + if self.closed: + raise RuntimeError("buffer pool is closed") + size = max(nbytes, 1) + if size > self.capacity: + raise RuntimeError("requested buffer size exceeds pool capacity") + if self.max_regions is not None and len(self.active) >= self.max_regions: + raise RuntimeError("buffer pool is exhausted") + offset = self._first_fit(size) + if offset is None: + raise RuntimeError("buffer pool is exhausted") + self.active[offset] = size + self.acquired.append(nbytes) + self.offsets.append(offset) + self.peak_bytes = max(self.peak_bytes, sum(self.active.values())) + self.peak_regions = max(self.peak_regions, len(self.active)) + return FakeBuffer(self, offset, nbytes) + + def _first_fit(self, size): + cursor = 0 + for offset in sorted(self.active): + if offset - cursor >= size: + return cursor + cursor = -(-(offset + self.active[offset]) // ALIGN) * ALIGN + return cursor if self.capacity - cursor >= size else None + + def _give_back(self, offset): + with self._lock: + del self.active[offset] + self.returned += 1 + + def close(self): + with self._lock: + if self.active: + raise RuntimeError("cannot close buffer pool with active leases") + self.closed = True + + +class PoolFactory: + """Stands in for mooncake.store.BufferPool, recording how the client configured it.""" + + def __init__(self, capacity=None): + self._capacity = capacity + self.kwargs: dict | None = None + self.pool: FakePool | None = None + + def __call__(self, store, max_bytes=0, max_regions=None): + self.kwargs = {"max_bytes": max_bytes, "max_regions": max_regions} + self.pool = FakePool(max_bytes if self._capacity is None else self._capacity, max_regions) + return self.pool + + +@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, **extra): + config = { + "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, + } + config.update(extra) + return mcc.MooncakeStoreClient(config) + + +def install_pool(monkeypatch, capacity=None): + """Make the client believe mooncake provides a buffer pool, and hand it a fake one.""" + factory = PoolFactory(capacity) + monkeypatch.setattr(mcc, "MOONCAKE_BUFFER_POOL_IMPORTED", True) + monkeypatch.setattr(mcc, "BufferPool", factory, raising=False) + return factory + + +def make_client_and_pool(monkeypatch, local_buffer_size=1 << 30, capacity=None, **extra): + factory = install_pool(monkeypatch, capacity) + client = make_client(local_buffer_size, **extra) + return client, factory.pool + + +def read(client, keys, shapes, dtypes): + tensors, indexes = client._get_tensors_thread_worker(keys, shapes, dtypes, list(range(len(keys)))) + assert indexes == list(range(len(keys))) + return tensors + + +def read_all(client): + return read(client, KEYS, SHAPES, DTYPES) + + +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 test_reads_land_in_pooled_buffer(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + + tensors = read_all(client) + + # Mixed dtypes and a scalar all round-trip through the staged copy-out. + assert_payloads(tensors) + # The whole point: no registration on the data path, and no buffer left on loan. + assert store.registered == [] + assert len(pool.acquired) == 1 and pool.returned == 1 and pool.active == {} + + +def test_pool_budget_is_capped_by_the_registered_local_buffer(store, monkeypatch): + factory = install_pool(monkeypatch) + client = make_client(local_buffer_size=1 << 20) + + # max_bytes=0 would let mooncake register fresh memory per lease once the local + # buffer is full, which is exactly the cost this path removes. + assert factory.kwargs == {"max_bytes": 1 << 20, "max_regions": mcc.MAX_BATCH_WORKER_THREADS} + # Half of an equal per-thread share, leaving the store room for its own staging. + assert client._lease_bytes == (1 << 20) // (2 * mcc.MAX_BATCH_WORKER_THREADS) + + +def test_batch_larger_than_share_is_read_in_rounds(store, monkeypatch): + # Small local buffer: each thread's share holds only part of the batch. + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + + 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.returned == len(pool.acquired) + assert store.registered == [] + + +def test_tensor_larger_than_share_registers_its_own_buffer(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + keys, shapes, dtypes = ["big", "small"], [(256,), (4,)], [torch.float32, torch.float32] + payloads = [torch.arange(256, dtype=torch.float32), torch.arange(4, dtype=torch.float32)] + store.objects = {k: bytes(t.numpy().tobytes()) for k, t in zip(keys, payloads, strict=True)} + + tensors = read(client, keys, shapes, dtypes) + + for got, want in zip(tensors, payloads, strict=True): + assert torch.equal(got, want) + # The oversized tensor must not lease the headroom the other readers need; the + # small one still takes the pooled path. + assert pool.acquired and all(nbytes <= client._lease_bytes for nbytes in pool.acquired) + assert store.registered and len(store.unregistered) == len(store.registered) + assert pool.active == {} + + +def test_falls_back_to_own_buffers_when_pool_is_exhausted(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + # Mooncake's own staging holds the arena, so acquire() raises instead of returning None. + held = pool.acquire(pool.capacity) + + tensors = read_all(client) + + assert_payloads(tensors) + assert pool.acquired == [pool.capacity] + assert store.registered and len(store.unregistered) == len(store.registered) + held.release() + + +def test_exhaustion_warns_once(store, monkeypatch, caplog): + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + held = pool.acquire(pool.capacity) + + with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): + read_all(client) + read_all(client) + + # Every group of every read falls back here; the hot path must not flood the log. + fallbacks = [r for r in caplog.records if "falling back to registering" in r.message] + assert len(fallbacks) == 1 + held.release() + + +def test_registers_receive_regions_without_pool_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. + client, pool = make_client_and_pool(monkeypatch) + n = 8 + 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 = read(client, keys, [(16,)] * n, [torch.float32] * n) + + for got, want in zip(tensors, payloads, strict=True): + assert torch.equal(got, want) + assert pool.acquired and pool.returned == len(pool.acquired) + assert store.registered == [] + + +def test_returned_tensors_survive_buffer_reuse(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + + first = read_all(client) + rounds = len(pool.acquired) + second = read_all(client) + + # The second read gets the same arena offsets back, so the first read's tensors + # prove the copy-out really handed ownership to the caller. + assert pool.offsets[rounds:] == pool.offsets[:rounds] + assert pool.returned == len(pool.acquired) + assert_payloads(first) + assert_payloads(second) + + +def test_failed_transfer_returns_the_buffer_and_keeps_the_cause(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + store.transfer_error = RuntimeError("transport is down") + + with pytest.raises(RuntimeError, match="transport is down"): + read_all(client) + + assert pool.active == {} and pool.returned == len(pool.acquired) + + +def test_failed_copy_out_returns_the_buffer_and_keeps_the_cause(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + + def boom(self, source, *args, **kwargs): + # Drop the staged view the caller passed in: a real copy_ failure comes from C++ + # and leaves no Python frame holding it, so only the client's own reference counts. + del source + raise RuntimeError("copy-out failed") + + monkeypatch.setattr(torch.Tensor, "copy_", boom) + + # A staged view left alive by the failing copy would make release() raise and + # replace the real cause, so the view has to be dropped on the error path too. + with pytest.raises(RuntimeError, match="copy-out failed"): + read_all(client) + + assert pool.active == {} and pool.returned == len(pool.acquired) + + +def test_concurrent_readers_stay_within_the_local_buffer(store, monkeypatch): + local_buffer_size = 8192 + readers = 8 + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=local_buffer_size) + barrier = threading.Barrier(readers) + # Hold every reader inside its transfer, so all of them need memory at the same time. + store.on_transfer = lambda: barrier.wait(timeout=30) + results: list[list[torch.Tensor]] = [None] * readers # type: ignore[list-item] + failures: list[Exception] = [] + + def reader(slot): + try: + results[slot] = read_all(client) + except Exception as e: # reported after the join below + failures.append(e) + + threads = [threading.Thread(target=reader, args=(i,)) for i in range(readers)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not failures + for tensors in results: + assert_payloads(tensors) + # Readers past the region cap fall back instead of eating into the store's headroom. + assert pool.peak_regions <= mcc.MAX_BATCH_WORKER_THREADS + assert pool.peak_bytes <= local_buffer_size // 2 + assert pool.active == {} + + +@pytest.mark.parametrize( + ("config", "warns"), + [ + ({"protocol": "rdma"}, True), + ({"protocol": "tcp"}, False), + ({"protocol": "rdma", "use_gdr": True}, False), + ], +) +def test_missing_pool_support_warns_only_for_cpu_rdma_reads(store, monkeypatch, caplog, config, warns): + monkeypatch.setattr(mcc, "MOONCAKE_BUFFER_POOL_IMPORTED", False) + # GDR staging defers cudaMalloc, so it can stand in without a CUDA device here. + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True) + + with caplog.at_level(logging.WARNING, logger=LOGGER_NAME): + make_client(**config) + + assert any("BufferPool is unavailable" in r.message for r in caplog.records) is warns + + +def test_close_returns_the_pool_before_the_store(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + + read_all(client) + client.close() + + assert pool.closed and store.closed and client._buffer_pool is None + + +def test_close_keeps_the_store_open_while_a_buffer_is_out(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + held = pool.acquire(1024) + + # The store owns the memory the buffer points into, so it must outlive the pool. + with pytest.raises(RuntimeError, match="while receive buffers are leased"): + client.close() + + assert not store.closed and client._buffer_pool is pool + held.release() + client.close() + assert pool.closed and store.closed diff --git a/tests/test_mooncake_client_lease.py b/tests/test_mooncake_client_lease.py deleted file mode 100644 index 4d27c344..00000000 --- a/tests/test_mooncake_client_lease.py +++ /dev/null @@ -1,206 +0,0 @@ -# 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 0fd85b2b..bd0b4076 100644 --- a/transfer_queue/storage/clients/mooncake_client.py +++ b/transfer_queue/storage/clients/mooncake_client.py @@ -54,14 +54,15 @@ 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. +def _copy_buffer_content_into_tensors(lease, targets: list[Tensor], offsets: list[int]) -> None: + """Copy each staged region out of the pooled 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. + fall back to a per-tensor copy. Views are dropped even when a copy raises: the + pool refuses to return a buffer while an exported view of it is still alive, and + an exception would otherwise keep the view alive in its traceback frame. """ t0 = targets[0] numel = t0.numel() @@ -83,17 +84,21 @@ def _copy_lease_into_tensors(lease, targets: list[Tensor], offsets: list[int]) - 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 + src = torch.frombuffer(lease.buffer, dtype=t0.dtype, count=(k - 1) * stride_elems + numel).as_strided( + (k, numel), (stride_elems, 1) + ) + try: + t0.as_strided((k, numel), (numel, 1), t0.storage_offset()).copy_(src) + finally: + 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 + try: + target.copy_(staged.view(target.shape)) + finally: + del staged @StorageClientFactory.register("MooncakeStoreClient") @@ -183,20 +188,22 @@ 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: + self._buffer_pool = ( + BufferPool(self._store, max_bytes=self.local_buffer_size, max_regions=MAX_BATCH_WORKER_THREADS) + if MOONCAKE_BUFFER_POOL_IMPORTED + else None + ) + self._buffer_pool_exhausted_logged = False + # Only CPU RDMA reads pay for registration; TCP and GDR reads do not benefit here. + if self._buffer_pool is None and self.protocol == "rdma" and self._gdr_staging 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." + "unregisters its own receive buffer. Upgrade to mooncake-transfer-engine >= 0.3.12 " + "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 + # Claim at most half the registered buffer: mooncake stages its own reads from the + # same region, so splitting all of it across the reader threads would starve those reads. + self._lease_bytes = self.local_buffer_size // (2 * MAX_BATCH_WORKER_THREADS) def put(self, keys: list[str], values: list[Any]) -> list[dict | None]: """Stores multiple key-value pairs to MooncakeStore. @@ -473,17 +480,17 @@ def _get_tensors_thread_worker( ) if self._buffer_pool is None: - self._read_into_own_buffers(batch_keys, batch_buffer_ptrs, batch_nbytes, region_ptrs, region_sizes) + self._read_into_registered_tensors(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) + self._read_via_buffer_pool(group, batch_keys, batch_buffer_ptrs, batch_nbytes, batch_buffer_tensors) return batch_buffer_tensors, indexes - def _read_into_own_buffers( + def _read_into_registered_tensors( 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.""" @@ -493,7 +500,7 @@ def _read_into_own_buffers( finally: self._unregister_all_buffers(region_ptrs) - def _read_group_via_lease( + def _read_via_buffer_pool( self, group: list[int], batch_keys: list[str], @@ -501,40 +508,56 @@ def _read_group_via_lease( batch_nbytes: list[int], batch_tensors: list[Tensor], ) -> None: - """Read one group of keys into leased memory, then copy into the caller's tensors. + """Read one group of keys into a pooled buffer, then copy into the caller's tensors. - The copy is what lets the lease return to the pool immediately, keeping the + The copy is what lets the buffer 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) + # split_by_bytes() gives a tensor larger than the share its own group without + # shrinking it, so the share is re-checked here rather than letting one reader + # lease the headroom the other readers and mooncake's own staging need. + lease = self._acquire_buffer(total) if total <= self._lease_bytes else None 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) + self._read_into_registered_tensors(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: + _copy_buffer_content_into_tensors(lease, [batch_tensors[i] for i in group], offsets) + except Exception: + # A failed read can leave the buffer unreturnable; keep the original cause. + try: + lease.release() + except Exception as release_error: + logger.warning(f"Returning the pooled receive buffer failed: {release_error}") + raise + else: lease.release() - def _acquire_lease(self, nbytes: int): + def _acquire_buffer(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. + Never block: waiting for a busy pool only delays a read that can register its own + receive buffer instead. Exhaustion is reported as an exception by mooncake. """ 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.") + # warn once and keep the details at debug level. + logger.debug(f"Leasing {nbytes} B of pre-registered memory failed ({e}); registering own buffer.") + if not self._buffer_pool_exhausted_logged: + self._buffer_pool_exhausted_logged = True + logger.warning( + "Cannot lease pre-registered memory from MooncakeStore, falling back to registering " + "receive buffers per read. Raise local_buffer_size to keep reads on the fast path." + ) return None def _get_tensors_gdr( @@ -655,10 +678,17 @@ def clear(self, keys: list[str], custom_backend_meta: list[Any] | None = None) - logger.error(f"remove failed for key `{actual_keys[i]}` with error code: {ret}") def close(self): - """Closes MooncakeStore.""" - # Release the leased regions before the store they belong to goes away. + """Closes MooncakeStore. + + Callers must have finished their reads first: the pool refuses to close while a + buffer is still leased, and the store owns the memory those buffers point into. + """ if self._buffer_pool is not None: - self._buffer_pool.close() + try: + self._buffer_pool.close() + except Exception as e: + # Keep the store open; its local buffer still backs the leased regions. + raise RuntimeError("Cannot close MooncakeStore while receive buffers are leased") from e self._buffer_pool = None if self._gdr_staging is not None: self._gdr_staging.close(self._store)