From 7a5245b2ddb277a8859801a244873ed3ce8c9306 Mon Sep 17 00:00:00 2001 From: chenjr97 Date: Wed, 11 Nov 2026 15:09:24 +0800 Subject: [PATCH] [perf] Reuse Mooncake buffer pool for CPU puts --- tests/test_mooncake_buffer_pool.py | 203 +++++++++++++++--- .../storage/clients/mooncake_client.py | 122 +++++++---- 2 files changed, 254 insertions(+), 71 deletions(-) diff --git a/tests/test_mooncake_buffer_pool.py b/tests/test_mooncake_buffer_pool.py index 400360ac..e142b291 100644 --- a/tests/test_mooncake_buffer_pool.py +++ b/tests/test_mooncake_buffer_pool.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Pooled receive buffers on the MooncakeStore tensor read path. +"""Pooled receive and send buffers on the MooncakeStore tensor read and put paths. 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 @@ -60,6 +60,7 @@ def __init__(self): self.closed = False self.transfer_error: Exception | None = None self.on_transfer = None + self.put_calls: list[tuple[list[str], list[int], list[int]]] = [] def setup(self, *args): return 0 @@ -81,6 +82,16 @@ def batch_get_into(self, keys, ptrs, sizes): ctypes.memmove(ptr, self.objects[key], size) return list(sizes) + def batch_upsert_from(self, keys, ptrs, sizes, config): + self.put_calls.append((list(keys), list(ptrs), list(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): + self.objects[key] = ctypes.string_at(ptr, size) + return [0] * len(keys) + def close(self): self.closed = True @@ -389,64 +400,79 @@ def test_returned_tensors_survive_buffer_reuse(store, monkeypatch): assert_payloads(second) -def test_failed_transfer_returns_the_buffer_and_keeps_the_cause(store, monkeypatch): +@pytest.mark.parametrize("operation", ["get", "put"]) +def test_failed_transfer_returns_the_buffer_and_keeps_the_cause(store, monkeypatch, operation): 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) + with pytest.raises(RuntimeError, match="transport is down") as exc: + if operation == "get": + read_all(client) + else: + client._put_tensors_thread_worker(KEYS, expected_tensors()) - assert pool.active == {} and pool.returned == len(pool.acquired) + assert exc.value is store.transfer_error + assert pool.active == {} and pool.returned == len(pool.acquired) == 1 -def test_failed_copy_out_returns_the_buffer_and_keeps_the_cause(store, monkeypatch): +@pytest.mark.parametrize("operation", ["get", "put"]) +def test_failed_copy_returns_the_buffer_and_keeps_the_cause(store, monkeypatch, operation): client, pool = make_client_and_pool(monkeypatch) + error = RuntimeError("copy failed") 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") + # A real copy_ failure leaves no Python mock frame holding either tensor. + # Drop both references so the mock does not keep pooled memory alive. + del self, source + raise error 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) + with pytest.raises(RuntimeError, match="copy failed") as exc: + if operation == "get": + read_all(client) + else: + client._put_tensors_thread_worker(KEYS, expected_tensors()) - assert pool.active == {} and pool.returned == len(pool.acquired) + assert exc.value is error + assert pool.active == {} and pool.returned == len(pool.acquired) == 1 -def test_concurrent_readers_stay_within_the_local_buffer(store, monkeypatch): +@pytest.mark.parametrize("mixed", [False, True], ids=["get_only", "mixed"]) +def test_concurrent_transfers_stay_within_the_local_buffer(store, monkeypatch, mixed): local_buffer_size = 8192 - readers = 8 + workers = 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. + barrier = threading.Barrier(workers) + # Hold every worker 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): + def transfer(slot): try: - results[slot] = read_all(client) + if mixed and slot % 2: + client._put_tensors_thread_worker([f"put{slot}/{key}" for key in KEYS], expected_tensors()) + else: + assert_payloads(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)] + threads = [threading.Thread(target=transfer, args=(i,)) for i in range(workers)] for t in threads: t.start() for t in threads: - t.join() + t.join(timeout=35) + assert not any(t.is_alive() for t in threads) 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. + if mixed: + for slot in range(1, workers, 2): + for key, tensor in zip(KEYS, expected_tensors(), strict=True): + assert store.objects[f"put{slot}/{key}"] == tensor.numpy().tobytes() + # Workers 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 == {} + assert pool.active == {} and pool.returned == len(pool.acquired) @pytest.mark.parametrize( @@ -507,3 +533,124 @@ def test_empty_tensor_reads_like_the_register_path(store, monkeypatch): for got, want in zip(pooled, registered, strict=True): assert got.shape == want.shape and torch.equal(got, want) assert pool.returned == len(pool.acquired) + + +def test_public_put_get_and_overwrite_mixed_values(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + keys = ["matrix", "metadata", "scalar", "tokens"] + matrix = torch.arange(12, dtype=torch.float32).reshape(3, 4).t() + assert not matrix.is_contiguous() + values = [matrix, {"step": 1, "labels": ["a", "b"]}, torch.tensor(17), torch.arange(5, dtype=torch.int16)] + shapes = [tuple(matrix.shape), None, (), (5,)] + dtypes = [torch.float32, None, torch.int64, torch.int16] + + meta = client.put(keys, values) + + assert meta[0] is None and meta[2:] == [None, None] + assert meta[1] == {"packed_size": len(store.objects["metadata"])} + assert store.registered == [] and store.unregistered == [] + first = client.get(keys, shapes, dtypes, meta) + for index in [0, 2, 3]: + assert first[index].dtype == values[index].dtype + assert first[index].shape == values[index].shape + assert torch.equal(first[index], values[index]) + assert first[1] == values[1] + + updated = [matrix + 10, {"step": 2, "labels": []}, torch.tensor(-3), values[3] + 4] + store.registered.clear() + store.unregistered.clear() + meta = client.put(keys, updated) + + assert store.registered == [] and store.unregistered == [] + for _, ptrs, sizes in store.put_calls: + for ptr, size in zip(ptrs, sizes, strict=True): + assert pool.arena.data_ptr() <= ptr < ptr + size <= pool.arena.data_ptr() + pool.capacity + second = client.get(keys, shapes, dtypes, meta) + for index in [0, 2, 3]: + assert torch.equal(second[index], updated[index]) + # Reusing pooled memory must not overwrite tensors returned by an earlier read. + assert torch.equal(first[index], values[index]) + assert second[1] == updated[1] + assert pool.active == {} and pool.returned == len(pool.acquired) + + +@pytest.mark.parametrize("fallback", ["missing_pool", "tcp", "exhausted"]) +def test_put_falls_back_to_registered_buffers(store, monkeypatch, fallback): + held = None + if fallback == "missing_pool": + monkeypatch.setattr(mcc, "MOONCAKE_BUFFER_POOL_IMPORTED", False) + client = make_client() + pool = None + else: + client, pool = make_client_and_pool(monkeypatch, protocol="tcp" if fallback == "tcp" else "rdma") + if fallback == "exhausted": + held = pool.acquire(pool.capacity) + tensor = torch.arange(12, dtype=torch.int16).reshape(3, 4).t() + assert not tensor.is_contiguous() + values = [tensor, {"fallback": fallback}] + + try: + meta = client.put(["tensor", "object"], values) + + assert store.registered and len(store.unregistered) == len(store.registered) + assert store.objects["tensor"] == values[0].numpy().tobytes() + assert meta == [None, {"packed_size": len(store.objects["object"])}] + result = client.get(["tensor", "object"], [(4, 3), None], [torch.int16, None], meta) + assert torch.equal(result[0], tensor) and result[1] == values[1] + if pool is not None: + assert pool.acquired == [pool.capacity] + finally: + if held is not None: + held.release() + + +def test_put_splits_groups_and_registers_only_oversized_values(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + keys = ["big", "small0", "small1", "small2", "object"] + values = [ + torch.arange(256, dtype=torch.float32), + torch.arange(4), + torch.arange(5), + torch.arange(6), + {"text": "x" * 1024}, + ] + + meta = client.put(keys, values) + + assert len(pool.acquired) == 2 + assert all(nbytes <= client._pool_share_bytes for nbytes in pool.acquired) + assert store.registered and len(store.registered) == len(store.unregistered) + for key, value in zip(keys[:4], values[:4], strict=True): + assert store.objects[key] == value.numpy().tobytes() + assert meta[4] == {"packed_size": len(store.objects["object"])} + assert pool.active == {} and pool.returned == len(pool.acquired) + + +def test_put_repeated_key_preserves_input_order(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch, local_buffer_size=4096) + values = [torch.arange(256, dtype=torch.float32), torch.arange(4, dtype=torch.float32)] + + assert client.put(["repeated", "repeated"], values) == [None, None] + + assert store.objects["repeated"] == values[-1].numpy().tobytes() + assert pool.active == {} + + +def test_put_keeps_the_same_lease_during_partial_retry(store, monkeypatch): + client, pool = make_client_and_pool(monkeypatch) + upsert = store.batch_upsert_from + + def retry_once(keys, ptrs, sizes, config): + assert len(pool.active) == 1 and pool.returned == 0 + results = upsert(keys, ptrs, sizes, config) + return [-1, 0] if len(store.put_calls) == 1 else results + + monkeypatch.setattr(store, "batch_upsert_from", retry_once) + tensors = [torch.arange(4), torch.arange(5)] + + client._put_tensors_thread_worker(["retry", "done"], tensors) + + initial, retry = store.put_calls + assert retry == (["retry"], [initial[1][0]], [initial[2][0]]) + assert store.objects["retry"] == tensors[0].numpy().tobytes() + assert pool.active == {} and pool.returned == len(pool.acquired) == 1 diff --git a/transfer_queue/storage/clients/mooncake_client.py b/transfer_queue/storage/clients/mooncake_client.py index e765f3c8..64f24921 100644 --- a/transfer_queue/storage/clients/mooncake_client.py +++ b/transfer_queue/storage/clients/mooncake_client.py @@ -82,27 +82,27 @@ def _resolve_npu_device() -> int | None: RETRY_DELAY_SECONDS = 1.0 -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. +def _copy_between_buffer_and_tensors(lease, tensors: list[Tensor], offsets: list[int], *, to_buffer: bool) -> None: + """Copy between the pooled buffer and tensors; ``to_buffer=True`` copies into the pool. Views are dropped even when a copy raises: the pool refuses to return a buffer while an exported view of it is alive, and an exception would otherwise keep it in its frame. """ # torch.frombuffer rejects count=0, and an empty tensor has nothing to copy. - staged = [(t, off) for t, off in zip(targets, offsets, strict=True) if t.numel()] + staged = [(t, off) for t, off in zip(tensors, offsets, strict=True) if t.numel()] if not staged: return - targets = [t for t, _ in staged] + tensors = [t for t, _ in staged] offsets = [off for _, off in staged] - t0 = targets[0] + t0 = tensors[0] numel = t0.numel() # A uniform, tightly packed group copies in one strided pass; ragged groups copy per tensor. - uniform = len(targets) > 1 + uniform = len(tensors) > 1 if uniform: base_off = t0.storage_offset() base_ptr = t0.untyped_storage().data_ptr() - for j, t in enumerate(targets): + for j, t in enumerate(tensors): if ( t.dtype != t0.dtype or t.shape != t0.shape @@ -114,21 +114,27 @@ def _copy_buffer_content_into_tensors(lease, targets: list[Tensor], offsets: lis break if uniform: - k = len(targets) + k = len(tensors) 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( + staged = 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) + if to_buffer: + staged.copy_(t0.as_strided((k, numel), (numel, 1), t0.storage_offset())) + else: + t0.as_strided((k, numel), (numel, 1), t0.storage_offset()).copy_(staged) finally: - del src + del staged return - for target, off in zip(targets, offsets, strict=True): - staged = torch.frombuffer(lease.buffer, dtype=target.dtype, count=target.numel(), offset=off) + for tensor, off in zip(tensors, offsets, strict=True): + staged = torch.frombuffer(lease.buffer, dtype=tensor.dtype, count=tensor.numel(), offset=off) try: - target.copy_(staged.view(target.shape)) + if to_buffer: + staged.view(tensor.shape).copy_(tensor) + else: + tensor.copy_(staged.view(tensor.shape)) finally: del staged @@ -223,8 +229,8 @@ def __init__(self, config: dict[str, Any]): if ret != 0: raise RuntimeError(f"Mooncake store setup failed with error code: {ret}") - # Only RDMA reads pay for register_buffer (a kernel op far slower than the transfer it - # enables); TCP reads would pay the extra copy-out for no gain, so leave them unpooled. + # Only RDMA transfers pay for register_buffer (a kernel op far slower than the transfer it + # enables); TCP transfers would pay the extra staging copy for no gain, so leave them unpooled. # See https://github.com/Ascend/TransferQueue/issues/169 use_pool = MOONCAKE_BUFFER_POOL_IMPORTED and self.protocol == "rdma" self._buffer_pool = ( @@ -235,12 +241,12 @@ def __init__(self, config: dict[str, Any]): self._pool_fallback_warned = False if not MOONCAKE_BUFFER_POOL_IMPORTED 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 to mooncake-transfer-engine >= 0.3.12 " + "mooncake.store.BufferPool is unavailable, so every tensor read and put registers and " + "unregisters its own buffer. Upgrade to mooncake-transfer-engine >= 0.3.12 " "to lease pre-registered buffers instead." ) # 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. + # same region, so splitting all of it across the read and put workers would starve those reads. self._pool_share_bytes = self.local_buffer_size // (2 * MAX_BATCH_WORKER_THREADS) def _bind_npu_device(self) -> None: @@ -377,7 +383,17 @@ def _put_tensors_gdr(self, batch_keys: list[str], batch_tensors: list[Tensor]) - def _put_tensors_thread_worker(self, batch_keys: list[str], batch_tensors: list[Tensor]) -> None: """Worker thread for putting tensors via the CPU RDMA path.""" - batch_ptrs, batch_sizes, _ = self._preprocess_tensors_for_put(batch_tensors) + batch_ptrs, batch_sizes, batch_tensors = self._preprocess_tensors_for_put(batch_tensors) + if self._buffer_pool is not None and len(set(batch_keys)) == len(batch_keys): + fallback: list[int] = [] + for group in split_by_bytes(batch_sizes, self._pool_share_bytes): + fallback += self._write_via_buffer_pool(group, batch_keys, batch_sizes, batch_tensors) + if not fallback: + return + batch_keys = [batch_keys[i] for i in fallback] + batch_ptrs = [batch_ptrs[i] for i in fallback] + batch_sizes = [batch_sizes[i] for i in fallback] + batch_ptr_reduced, batch_sizes_reduced = merge_contiguous_memory(batch_ptrs, batch_sizes) self._register_all_buffers(batch_ptr_reduced, batch_sizes_reduced) try: @@ -385,15 +401,42 @@ def _put_tensors_thread_worker(self, batch_keys: list[str], batch_tensors: list[ finally: self._unregister_all_buffers(batch_ptr_reduced) + def _write_via_buffer_pool( + self, + group: list[int], + batch_keys: list[str], + batch_nbytes: list[int], + batch_tensors: list[Tensor], + ) -> list[int]: + """Copy one group into a pooled buffer and upsert it from there. + + Returns the group unwritten when the pool cannot serve it, so the caller can register + those tensors in one batch. As on the read path, a group larger than one worker's share + is never leased. + """ + nbytes = [batch_nbytes[i] for i in group] + offsets, total = _aligned_offsets(nbytes) + lease = self._acquire_buffer(total) if total <= self._pool_share_bytes else None + if lease is None: + return group + + try: + _copy_between_buffer_and_tensors(lease, [batch_tensors[i] for i in group], offsets, to_buffer=True) + keys = [batch_keys[i] for i in group] + self._batch_upsert_with_retry(keys, [lease.ptr + off for off in offsets], nbytes) + except Exception: + try: + lease.release() + except Exception as release_error: + logger.warning(f"Returning the pooled send buffer failed: {release_error}") + raise + lease.release() + return [] + def _put_bytes_thread_worker(self, batch_keys: list[str], batch_values: list[Any]) -> list[int]: """Worker thread for putting batch of non-tensors to MooncakeStore.""" - # TODO: switch to a pre-registered buffer from MooncakeStore once such an API is available. - region_ptrs: list[int] = [] - region_sizes: list[int] = [] - def alloc(sizes: list[int]) -> list[Tensor]: - nonlocal region_ptrs, region_sizes # `batch_packed_sizes` are byte counts. With torch.uint8 (1 byte/element), # a 1-D shape of (N,) corresponds to exactly N bytes. We use # `allocate_empty_tensors` to get N uint8 views over a single contiguous, @@ -401,19 +444,13 @@ def alloc(sizes: list[int]) -> list[Tensor]: # consumers apply the actual dtype/shape interpretation when unpacking. dtypes = [torch.uint8] * len(sizes) shapes = [(s,) for s in sizes] - buffers, _, region_ptrs, region_sizes = allocate_empty_tensors(dtypes, shapes) + buffers, _, _, _ = allocate_empty_tensors(dtypes, shapes) return buffers buffers, batch_sizes = serial_utils.batch_encode_into( batch_values, alloc, num_workers=MAX_SERIAL_WORKER_THREADS ) - batch_ptrs = [cast(Tensor, b).data_ptr() for b in buffers] - - self._register_all_buffers(region_ptrs, region_sizes) - try: - self._batch_upsert_with_retry(batch_keys, batch_ptrs, batch_sizes) - finally: - self._unregister_all_buffers(region_ptrs) + self._put_tensors_thread_worker(batch_keys, cast(list[Tensor], buffers)) return batch_sizes @@ -544,7 +581,7 @@ def _get_tensors_thread_worker( fallback: list[int] = [] for group in split_by_bytes(batch_nbytes, self._pool_share_bytes): fallback += self._read_via_buffer_pool( - group, batch_keys, batch_buffer_ptrs, batch_nbytes, batch_buffer_tensors + group, batch_keys, batch_nbytes, batch_buffer_tensors ) if fallback: self._read_into_registered_tensors( @@ -568,15 +605,14 @@ def _read_via_buffer_pool( self, group: list[int], batch_keys: list[str], - batch_ptrs: list[int], batch_nbytes: list[int], batch_tensors: list[Tensor], ) -> list[int]: """Read one group into a pooled buffer and copy into the caller's tensors. Returns the group unread when the pool cannot serve it, so the caller can register - those keys in one batch. A group larger than one reader's share is never leased: - that headroom belongs to the other readers and to mooncake's own staging. + those keys in one batch. A group larger than one worker's share is never leased: + that headroom belongs to the other workers and to mooncake's own staging. """ nbytes = [batch_nbytes[i] for i in group] offsets, total = _aligned_offsets(nbytes) @@ -587,7 +623,7 @@ def _read_via_buffer_pool( keys = [batch_keys[i] for i in group] try: self._batch_get_into_with_retry(keys, [lease.ptr + off for off in offsets], nbytes) - _copy_buffer_content_into_tensors(lease, [batch_tensors[i] for i in group], offsets) + _copy_between_buffer_and_tensors(lease, [batch_tensors[i] for i in group], offsets, to_buffer=False) except Exception: # A failed read can leave the buffer unreturnable; keep the original cause. try: @@ -601,8 +637,8 @@ def _read_via_buffer_pool( def _acquire_buffer(self, nbytes: int): """Lease ``nbytes`` of pre-registered memory, or None when the pool cannot serve it. - Never block: a busy pool means registering a receive buffer beats waiting. mooncake - reports both the per-reader cap and an over-capacity request as an exception. + Never block: a busy pool means registering a transfer buffer beats waiting. mooncake + reports both the per-worker cap and an over-capacity request as an exception. """ assert self._buffer_pool is not None try: @@ -612,8 +648,8 @@ def _acquire_buffer(self, nbytes: int): if not self._pool_fallback_warned: self._pool_fallback_warned = True logger.warning( - f"Falling back to registering receive buffers per read ({e}). Raise local_buffer_size " - "or lower read concurrency to keep reads on the pre-registered fast path." + f"Falling back to registering transfer buffers per call ({e}). Raise local_buffer_size " + "or lower read and put concurrency to stay on the pre-registered fast path." ) return None @@ -737,7 +773,7 @@ def clear(self, keys: list[str], custom_backend_meta: list[Any] | None = None) - def close(self): """Closes MooncakeStore. - Callers must have finished their reads first: the pool refuses to close while a + Callers must have finished their reads and puts 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: