From 961ee1277eef3f231a70091b42474ff7694cdc75 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Tue, 27 Oct 2026 08:23:02 +0000 Subject: [PATCH 01/10] [None][feat] Add the Mooncake store driver MooncakeDriver implements the StoreDriver protocol over a set-up MooncakeDistributedStore from SDK 0.3.13, zero copy from and into the lender's registered host slots. A put starts a put session per key, writes every row in one scatter (batch_put_from_multi_buffer_ranges), ends the sessions of rows written whole and revokes the rest; a get is one batch_get_into. The module docstring cites the 0.3.13 sources each rule was read from. - put and get run on the executor loop's thread: a few comparisons whatever the batch's size (one span per row, at most 2 GiB a row, inside one live registration), then the operation is recorded and queued whole. Admitting one costs the same for any number of rows or parts. - A fixed pool of worker threads, started with the driver and ending when it drains, cuts parts off the queued operations, at most rows_per_call rows and bytes_per_call bytes each, and makes each part's SDK calls. Gets go first; once GETS_BEFORE_PUT_IS_DUE gets have gone ahead of a waiting put, that put goes next once a put may run, and puts take at most workers - 1 threads, so a get never waits behind puts alone. - on_change is called once an operation is done or turned unprovable, outside the driver's lock. A cancel ends the rows no worker took, and a drain those of every operation; the running parts run to their end. A worker an exception stops faults its part with access unproven and the others go on; once none is left, the rows no worker took fault and new operations are rejected. - Access end differs by call. The put scatter waits, without a timeout, until every task it posted finished, so its return ends access whatever it answers. batch_get_into can answer TRANSFER_FAIL for a future that gave up while its transfer still runs, so a get ends access only when no row answers TRANSFER_FAIL. A raise or a wrong answer count leaves access unproven. An operation whose access cannot be proven keeps its region from being unregistered and makes drain answer False at once. - Put sessions are per key and client. The driver never starts a session for a key one of its puts holds (that row is NOT_READY without asking the SDK) and ends only rows written whole. A session it cannot see ended or revoked stays in one ledger: left for drain to revoke (an end whose outcome is unknown, a start or put that raised or miscounted, a failed revoke), or left to the master when the target may still write into the object (TRANSFER_FAIL, a scatter that raised or miscounted). A key whose session is left open is never put again: it is READY when the master shows its object complete. - Codes become row results here and nowhere else: OBJECT_NOT_FOUND is a MISS; a key with no complete replica, a key held by this driver's put and an uncommitted object are NOT_READY; a full pool is DECLINED; an existing object is READY once the master confirms it; a read whose lease expired is read once more, READY or FAULT; the rest are FAULTs with their codes. Skipped rows are counted by reason in stats() and named in the part's warning. - Capabilities: absent_is_miss, host memory only, one span per row, cancel, no probe and no row limit of its own. close() closes the SDK client once a drain returned True and every region is unregistered. fake_mooncake.py implements only the SDK calls the driver may make, with 0.3.13's answers over real bytes; test_mooncake_driver.py pins the code mapping, sessions, access-end rules, scheduling, drain and close, a backend round trip, and that fetch and publish build no store key or RowIO per row and start no thread on the calling thread. The shared driver_conformance checks run over the fake SDK with every scenario it can stage. An exception escaping a thread fails the test; the tests that end a worker thread on purpose assert it ended by SystemExit. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../_torch/disaggregation/blob/README.md | 3 + .../_torch/disaggregation/blob/__init__.py | 1 + .../_torch/disaggregation/blob/mooncake.py | 1038 ++++++++++++ .../disaggregation/blob/fake_mooncake.py | 299 ++++ .../blob/test_mooncake_driver.py | 1417 +++++++++++++++++ 5 files changed, 2758 insertions(+) create mode 100644 tensorrt_llm/_torch/disaggregation/blob/mooncake.py create mode 100644 tests/unittest/_torch/disaggregation/blob/fake_mooncake.py create mode 100644 tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py diff --git a/tensorrt_llm/_torch/disaggregation/blob/README.md b/tensorrt_llm/_torch/disaggregation/blob/README.md index 223f1a0bc9c9..612437cc5695 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/README.md +++ b/tensorrt_llm/_torch/disaggregation/blob/README.md @@ -32,6 +32,7 @@ store SDK | `backend.py` | `BlobBackend`, `BlobAttempt`, `BlobConfig`, `DriverCallLimit`, `RegionBinding`, `ConstructionRejected` | | `store.py` | the driver boundary: `StoreDriver`, the optional `RegistersRegionsTogether`, `RowBatch`, `RowIO`, `OnChange`, `RowStatus`, `DriverCapabilities`, `DriverRejected`, `RegistrationError` | | `keys.py` | the store key encoding | +| `mooncake.py` | `MooncakeDriver`, a `StoreDriver` over a set-up Mooncake 0.3.13 SDK client | | `../testing/` | what a driver's own CI imports from the wheel: the conformance suite, the in-memory store and drivers, the broken drivers that prove every check bites | ## Running the tests @@ -52,6 +53,8 @@ All of it is CPU-only. The files, top down: rule it guards. - `test_blob_format.py`: keys, spans, layouts, row batches, the object trailer. - `test_testing_isolation.py`: no engine module imports the test tooling. +- `test_mooncake_driver.py`: the Mooncake driver over a fake SDK (`fake_mooncake.py`) that + answers as 0.3.13 does, the conformance suite included. ## Writing a driver diff --git a/tensorrt_llm/_torch/disaggregation/blob/__init__.py b/tensorrt_llm/_torch/disaggregation/blob/__init__.py index 80de5d3d1a62..7809fd32d115 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/blob/__init__.py @@ -12,6 +12,7 @@ - ``store`` and ``keys``: the private driver boundary, which shares only ``MemoryKind`` with the public contract, and the store key encoding. How an object is laid out in a store is its driver's private format. +- ``mooncake``: ``MooncakeDriver``, a ``StoreDriver`` over a set-up Mooncake SDK client. No module imports torch, ``tensorrt_llm._torch.pyexecutor`` or a store SDK at import time. ``README.md`` next to this file says how to write a driver and how to assemble a backend. diff --git a/tensorrt_llm/_torch/disaggregation/blob/mooncake.py b/tensorrt_llm/_torch/disaggregation/blob/mooncake.py new file mode 100644 index 000000000000..94351d1a222f --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/blob/mooncake.py @@ -0,0 +1,1038 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A store driver over Mooncake's distributed store, pinned to SDK 0.3.13. + +Each row is one object holding exactly the row's payload, written from and read into the lender's +registered host slots without a copy. ``put`` and ``get`` run on the executor loop's thread: they +check a batch with a few comparisons, whatever its size, and record and queue the operation. +``workers`` threads, started with the driver and running until it drains, take the queued rows a +part at a time, at most ``rows_per_call`` rows and ``bytes_per_call`` payload bytes, and make each +part's SDK calls, which block their thread until the store answers. Gets go first, as they sit on +the time-to-first-token path, but once ``GETS_BEFORE_PUT_IS_DUE`` gets have gone ahead of a +waiting put, that put goes next once a put may run. Puts occupy at most ``workers - 1`` threads (at +least one), so with two workers or more a get never waits behind puts alone. Every SDK return code +is translated here and nowhere else. + +A part's rows are READY unless an SDK answer says otherwise: a worker compares a call's answers +with what an all-OK call answers, and looks at rows one by one only when they differ. + +Put, per part: ``batch_put_session_start`` per key (the master reserves the object), one +``batch_put_from_multi_buffer_ranges`` scatter for all its rows, then ``batch_put_session_end`` for +the rows written whole (their objects become readable) and ``batch_put_session_revoke`` for the +rest. Put uses sessions because this scatter has no timeout, so a failed row is revoked only once +nothing writes into its space; ``batch_put_from`` revokes a timed-out transfer at once, while its +RDMA writes may still land in space the master has given to another object. + +Get, per part: one ``batch_get_into``. A get session would add two SDK calls, and every call waits +for the GIL on return, which a busy executor thread makes cost milliseconds. + +When access ends. The rules below were read in the Mooncake sources at tag v0.3.13 (paths are in +that tree), for a wheel built without ``USE_TENT``, as the release wheels are +(``mooncake-common/common.cmake:174``, ``.github/workflows/_build-wheel.yaml:67-107``). They hold +only if whoever opens the SDK refuses ``MC_USE_TENT`` and ``MC_USE_TEV1``, which switch to the +newer engine at run time (``mooncake-transfer-engine/src/transfer_engine.cpp:393-410``); the SDK's +local hot cache, which could serve a torn read later; and TCP's legacy framing, which reports a +write before it lands and is also used toward a process of the pool that runs an SDK older than +0.3.12. The default replicate config is assumed: one memory replica, no NoF or SSD replica. + +- The put scatter waits, without a timeout, until every task it published is terminal, also when + submitting failed part way (``transfer_engine.cpp:987-993, 1236-1265, 1315-1319``). An RDMA task + is FAILED only once all its slices ended (``mooncake-transfer-engine/src/transport/ + rdma_transport/rdma_transport.cpp:1023-1048``), and a TCP v2 write ends, either way, only after + the body write released the caller's buffer (``mooncake-transfer-engine/src/transport/ + tcp_transport/tcp_transport_session_impl.h:818-905``). So the + scatter's return ends access to the slots whatever it answers, TRANSFER_FAIL included. +- Over TCP a TRANSFER_FAIL may also mean the initiator gave up on a slow target: no + acknowledgment within ``MC_TCP_STATUS_TIMEOUT_SEC`` (30 s by default; + ``tcp_transport_session_impl.h:394-401, 447-470``) or no progress within + ``MC_TCP_PROGRESS_TIMEOUT_SEC`` (``:404-414, 496-526``), while the target may still apply data + it received into the object's space (``:838-840``). A revoke frees that space at once + (``mooncake-store/src/master_service.cpp:4776, 4869``), so such a row's session is never revoked + by the driver; the master frees a put left unfinished only after + ``put_start_release_timeout_sec`` (600 s by default; ``master_service.cpp:4441-4450``). +- ``batch_get_into`` waits on one transfer future per key, and a future can give up with + TRANSFER_FAIL (after 60 s, on a status error, after a partly posted RDMA batch; + ``mooncake-store/src/transfer_task.cpp:746-905, 1239-1258``) while its transfer still runs; + every other answer comes after all futures finished or from a step that posted nothing. So a get + ends access only when no answer is TRANSFER_FAIL. +- In both, a raise or a wrong answer count leaves access unproven. Session, existence and end + calls touch only master and local metadata, never row memory. + +An operation can never prove its access ended once one of its parts returned from its SDK calls +with access not ended: it is ``access_unprovable``, and it keeps its region from being unregistered +and the driver from draining. Its other parts still run: each part's rows stand alone, and a cancel +or drain ends only the rows no worker took yet. +""" + +from __future__ import annotations + +import enum +import itertools +import threading +import time +from collections import Counter, deque +from collections.abc import Callable, Iterable, Sequence +from dataclasses import dataclass +from enum import IntEnum +from typing import Any, NamedTuple + +from tensorrt_llm.logger import logger + +from ..base.memory import MemoryKind +from .store import ( + DriverCapabilities, + DriverRejected, + Hint, + OnChange, + ProbeStatus, + RegionLayout, + RegistrationError, + RowBatch, + RowResult, + RowStatus, + Span, + StoreError, + faulted, +) + +__all__ = ["MooncakeDriver"] + + +class _Code(IntEnum): + """The SDK return codes the driver names (``mooncake-store/include/types.h``).""" + + OK = 0 + INTERNAL_ERROR = -1 + NO_AVAILABLE_HANDLE = -200 + INVALID_PARAMS = -600 + ILLEGAL_CLIENT = -601 + INVALID_WRITE = -700 + INVALID_REPLICA = -702 + REPLICA_IS_NOT_READY = -703 + OBJECT_NOT_FOUND = -704 + OBJECT_ALREADY_EXISTS = -705 + LEASE_EXPIRED = -707 + TRANSFER_FAIL = -800 + CHECKSUM_MISMATCH = -801 # with MOONCAKE_STORE_CHECKSUM: the bytes of a get landed, then failed + RPC_FAIL = -900 + RPC_TIMEOUT = -901 + # Not an SDK code: the driver's own answer when a metadata call raised or miscounted (logged). + METADATA_CALL_FAILED = -100_000 + + +OBJECT_EXISTS, OBJECT_ABSENT = 1, 0 # batch_is_exist answers +DEFAULT_WORKERS = 4 +DEFAULT_ROWS_PER_CALL = 256 +DEFAULT_BYTES_PER_CALL = 128 << 20 +# Gets taken while a put waits; then that put goes next, as soon as a put may run. +GETS_BEFORE_PUT_IS_DUE = 4 +MAX_ROW_BYTES = (1 << 31) - 1 # the put scatter answers each row's bytes as a 32-bit int +_CAPABILITIES = DriverCapabilities( + get=True, + put=True, + cancel=True, # ends the rows no worker took; the SDK cannot stop a running call + probe=False, + max_rows_per_op=None, # any number: the driver cuts an operation into parts itself + # The master answers OBJECT_NOT_FOUND for a key it does not hold; failures keep their own + # codes and stay FAULT. + absent_is_miss=True, + device_memory=False, # registers host memory only: the lender stages rows there + # One span per row: batch_get_into reads each object into one buffer; a get of several + # spans would need batch_get_into_multi_buffers. + multi_span=False, +) +# The results that carry no fault text, one shared object each. +_PENDING = RowResult(RowStatus.PENDING) +_READY = RowResult(RowStatus.READY) +_MISS = RowResult(RowStatus.MISS) +_NOT_READY = RowResult(RowStatus.NOT_READY) +_DECLINED = RowResult(RowStatus.DECLINED) +_CANCELLED = RowResult(RowStatus.CANCELLED) +_NO_WORKER_LEFT = faulted("no worker thread is left to run the row") +_instance_ids = itertools.count() # names each driver's worker threads + +# Why a row was skipped (nothing moved and nothing went wrong), and the result it gets; ``stats()`` +# counts skipped rows by both, and a part's warning names its first row that faulted or was +# skipped, with the skip's reason. +_SKIP_RESULTS = { + # get -703: no complete replica, as while the object is written or once its replica's + # endpoint went stale + "no_complete_replica": _NOT_READY, + # put: a put of this driver is writing the key now, or a session it left open holds the key and + # the master shows no complete object; no session is started + "held_by_own_put": _NOT_READY, + "uncommitted": _NOT_READY, # put start -705, and the master shows the object uncommitted + "pool_full": _DECLINED, # put start -200: no space for the object +} + +# Session-end answers after which the end may still reach the master, even if the retry is refused. +# Such a session is left for drain and its key is never put again: the master matches an end by key +# and client only, so a late end would mark complete this client's next object of the key. +_END_OUTCOME_UNKNOWN = (_Code.RPC_FAIL, _Code.RPC_TIMEOUT, _Code.METADATA_CALL_FAILED) +# Session-end answers retried once, as the master ends a session idempotently; an end call that +# raised or miscounted is not retried. +_END_RETRIED = (_Code.RPC_FAIL, _Code.RPC_TIMEOUT) + + +class _LeftOpen(enum.Enum): + """Who ends a put session this driver left open.""" + + FOR_DRAIN = "for_drain" # drain revokes it once no part runs + TO_MASTER = "to_master" # the target may still write into its object: the master frees it + + +@dataclass(eq=False) +class _Registration: + span: Span + live: bool = True + ops_accessing: int = 0 # operations whose access may touch it; unregister refuses it until 0 + + +@dataclass(eq=False) +class _Op: + """One ``put`` or ``get``; its fields change under the driver's lock.""" + + owner: MooncakeDriver + is_put: bool + rows: RowBatch + region: _Registration # the registration every row lies in + rows_per_part: int + results: list[RowResult] # per row; each part writes its slice when it ends + # Called once every row concluded (its access ends then too, unless unprovable) and once + # the operation turned unprovable. + on_change: OnChange | None + # Rows from here on wait in the operation's queue; once every row was taken, or ended + # untaken, it is ``len(rows)`` and the operation has left the queue. + first_untaken_row: int = 0 + parts_running: int = 0 + # A part returned with its access not proven ended: the region stays busy for good. + unprovable: bool = False + + @property + def all_taken(self) -> bool: + return self.first_untaken_row == len(self.rows) + + @property + def access_ended(self) -> bool: + return self.all_taken and self.parts_running == 0 and not self.unprovable + + +@dataclass(frozen=True, eq=False) +class _Part: + """Rows ``[start, stop)`` of ``op``: what a worker cut off the operation's queued rows when it + took them, and runs in one go.""" + + op: _Op + start: int + stop: int + + +class _Rows: + """Rows of one part as SDK calls take them: keys, addresses, and the size every row has.""" + + def __init__(self, keys: list[str], addresses: list[int], row_bytes: int) -> None: + self.keys = keys + self.addresses = addresses + self.row_bytes = row_bytes + # Also what a call answers when it moved every row whole. + self.sizes = [row_bytes] * len(keys) + + def only(self, keys: list[str]) -> _Rows: + """The rows of ``keys``, a subset of these in their order.""" + if len(keys) == len(self.keys): + return self + address_of = dict(zip(self.keys, self.addresses)) + return _Rows(keys, [address_of[key] for key in keys], self.row_bytes) + + +class _PartResults: + """A part's row results, recorded only for the rows that did not end READY, by key (a batch + never repeats a key), and why each skipped row was skipped (a key of ``_SKIP_RESULTS``).""" + + def __init__(self) -> None: + self.recorded: dict[str, RowResult] = {} + self.skip_reasons: dict[str, str] = {} + + @property + def all_ready(self) -> bool: + return not self.recorded + + def skip(self, key: str, reason: str) -> None: + self.recorded[key] = _SKIP_RESULTS[reason] + self.skip_reasons[key] = reason + + def fault(self, keys: Iterable[str], reason: str) -> None: + result = faulted(reason) + for key in keys: + self.recorded[key] = result + + def in_order(self, keys: list[str]) -> list[RowResult]: + """The results of ``keys``, READY where none was recorded.""" + if not self.recorded: + return [_READY] * len(keys) + return [self.recorded.get(key, _READY) for key in keys] + + +class _PartEnd(NamedTuple): + """How a part's SDK calls ended: its rows' results in order, the reasons its rows were skipped, + and whether their access ended.""" + + results: list[RowResult] + skip_reasons: tuple[str, ...] + access_ended: bool + + +class _WrongAnswerCount(Exception): + """The SDK answered for a different number of keys than it was given; its state is unknown.""" + + +class MooncakeDriver: + """A ``StoreDriver`` over one set-up ``mooncake.store.MooncakeDistributedStore``. It must be the + only driver over its SDK client: the SDK names a put session by its key, per client. + + A put of a part's rows: + + 1. Reserves each key that no put of this driver is writing and no session it left open holds. + A key another put is writing is NOT_READY at once. A key whose session it left open is + READY when the master shows its object complete, NOT_READY when it shows none, and FAULT + when the existence check fails; it is never put again. + 2. Starts a session per reserved key, writes every started row in one scatter, ends the + session of each row written whole (its object becomes readable) and revokes the rest. + 3. Leaves open each session it may have started and did not see ended or revoked: + + - to the master, never revoked here, since the target may still write into the object: a + scatter that raised or miscounted, a row answered TRANSFER_FAIL. The master frees them + after ``put_start_release_timeout_sec``. + - for ``drain`` to revoke once no part runs: any other, such as a start that raised or + miscounted, an end whose outcome is unknown (``_END_OUTCOME_UNKNOWN``), a revoke that + failed, or a put stopped by a raise. + + Args: + store: An SDK store whose ``setup`` returned 0. + workers: The most threads running SDK calls; each blocks while its call waits on the + store. + rows_per_call: The most rows of one part, so of each SDK call. + bytes_per_call: The most payload bytes of one part, so of each SDK call that moves rows; + a row larger than this goes in a part of its own. + """ + + def __init__( + self, + store: Any, + *, + workers: int = DEFAULT_WORKERS, + rows_per_call: int = DEFAULT_ROWS_PER_CALL, + bytes_per_call: int = DEFAULT_BYTES_PER_CALL, + ) -> None: + for name, value in ( + ("workers", workers), + ("rows_per_call", rows_per_call), + ("bytes_per_call", bytes_per_call), + ): + if not isinstance(value, int) or isinstance(value, bool) or value < 1: + raise ValueError(f"{name} must be a positive int, got {value!r}") + self._store = store + self._rows_per_call = rows_per_call + self._bytes_per_call = bytes_per_call + # Puts occupy at most this many threads, so with two workers or more a get never waits + # behind puts alone. + self._max_running_puts = max(1, workers - 1) + # Guards the fields below and wakes ``drain``; never held across an SDK call or an + # ``on_change``. Idle workers wait on ``_work_due``, which shares its lock. + mutex = threading.Lock() + self._lock = threading.Condition(mutex) + self._work_due = threading.Condition(mutex) + # Operations with rows no worker took yet, oldest first. + self._queued_gets: deque[_Op] = deque() + self._queued_puts: deque[_Op] = deque() + self._workers_alive = 0 # worker threads started and not yet exited + self._running_puts = 0 # put parts taken and not yet ended; at most _max_running_puts + # Get parts taken while a put waited, since a put part was taken or a put started waiting + # with none queued; at GETS_BEFORE_PUT_IS_DUE the put is due. + self._gets_ahead_of_puts = 0 + self._registrations: list[_Registration] = [] + # Set by the first drain, whatever it returns, or when a worker thread fails to start. + self._draining = False + # A drain returned True: no part ran, no worker thread was left and none was unprovable. + self._drained = False + self._closed = False + # Keys a running put part reserved, until that part ends: no two puts of this driver hold + # a session of one key at once, so a -600 is never this driver's own. + self._keys_being_put: set[str] = set() + # The put sessions this driver left open (class docstring, step 3), by key. No key is in + # both: a reserved key is left open only as its reservation ends, and is never reserved + # again. + self._sessions_left_open: dict[str, _LeftOpen] = {} + self._in_flight_op_count = 0 # operations with rows queued or parts running + self._unprovable_op_count = 0 # operations whose access can never be proven ended + self._skipped_rows: Counter[str] = Counter() # by reason + # Started here, on the builder's thread: starting a thread waits for it to run, which on + # the executor loop's thread would hand the GIL to every busy thread. If one cannot start, + # those already started exit, as after a drain. + name = f"mooncake-driver-{next(_instance_ids)}-worker" + with self._lock: + try: + for _ in range(workers): + threading.Thread(target=self._worker_loop, name=name, daemon=True).start() + self._workers_alive += 1 + except BaseException: + self._draining = True + self._work_due.notify_all() + raise + + @property + def capabilities(self) -> DriverCapabilities: + return _CAPABILITIES + + def register(self, address: int, nbytes: int, layout: RegionLayout) -> _Registration: + try: + span = Span(address, nbytes) + except ValueError as error: + raise DriverRejected(str(error)) from None + if layout.nbytes > nbytes: + raise DriverRejected(f"layout needs {layout.nbytes} bytes, region has {nbytes}") + if layout.memory_kind is not MemoryKind.HOST: + raise DriverRejected("this driver registers host memory only") + with self._lock: + self._require_open(DriverRejected) + overlapping = [r.span for r in self._registrations if r.live and r.span.overlaps(span)] + if overlapping: + raise DriverRejected(f"region {span} overlaps a live registration {overlapping[0]}") + registration = _Registration(span) + try: + code = self._store.register_buffer(address, nbytes) + except Exception as error: + self._add(registration) + raise RegistrationError(f"register_buffer raised {error!r}", registration) from error + # The SDK undoes a partial registration before it answers (RdmaTransport:: + # registerLocalMemoryInternal), so a non-OK answer leaves nothing registered; a rollback + # that itself fails is only logged by the SDK and cannot be seen here. + if code != _Code.OK: + raise RegistrationError(f"register_buffer returned {_describe(code)}") + self._add(registration) + return registration + + def unregister(self, handle: _Registration) -> None: + with self._lock: + registration = self._checked_registration(handle) + if not registration.live: + return + if registration.ops_accessing: + raise DriverRejected( + f"{registration.ops_accessing} operation(s) may still access the region" + ) + registration.live = False # admits no new rows while the SDK unregisters it + try: + code = self._store.unregister_buffer(registration.span.address) + except Exception as error: + failure = StoreError(f"unregister_buffer raised {error!r}") + else: + if code == _Code.OK: + return + failure = StoreError(f"unregister_buffer returned {_describe(code)}") + with self._lock: + registration.live = True # still registered: the call may be retried + raise failure + + def put(self, rows: RowBatch, *, on_change: OnChange | None = None) -> _Op: + return self._admit(rows, is_put=True, on_change=on_change) + + def get(self, rows: RowBatch, hint: Hint = None, *, on_change: OnChange | None = None) -> _Op: + if hint is not None: + raise DriverRejected("this driver reads no hints") + return self._admit(rows, is_put=False, on_change=on_change) + + def poll(self, op: _Op) -> tuple[RowResult, ...]: + with self._lock: + return tuple(self._checked_op(op).results) + + def access_ended(self, op: _Op) -> bool: + with self._lock: + return self._checked_op(op).access_ended + + def access_unprovable(self, op: _Op) -> bool: + with self._lock: + return self._checked_op(op).unprovable + + def cancel(self, op: _Op) -> None: + """Cancels the rows of ``op`` no worker took; the running parts run to their end.""" + with self._lock: + done = self._end_untaken_rows([self._checked_op(op)], _CANCELLED) + if done: + self._call_on_change(op) + + def probe( + self, keys: Sequence[str], budget_s: float, hint: Hint = None + ) -> tuple[ProbeStatus, ...]: + """Never called: ``capabilities.probe`` is false.""" + self._require_open() + return tuple(ProbeStatus.UNKNOWN for _ in keys) + + def drain(self, budget_s: float) -> bool: + """Refuses new work for good, cancels the rows no worker took, ends the worker threads, + and waits up to ``budget_s`` for the running parts to end and every worker thread to exit. + Returns False at once while an operation is unprovable. + + Each drain that returns True also tries once to revoke the sessions left for drain: no put + runs any more, so none is in use. They hold no lender memory, so they never change the + answer; one whose revocation fails stays for the next drain, or until the master discards + it. Sessions left to the master are never revoked here.""" + with self._lock: + self._require_open() + self._draining = True + done = self._end_untaken_rows(self._queued_ops(), _CANCELLED) # none starts now + self._work_due.notify_all() # idle workers exit + for op in done: + self._call_on_change(op) + with self._lock: + self._lock.wait_for( + lambda: self._unprovable_op_count + or not (self._in_flight_op_count or self._workers_alive), + min(max(0.0, budget_s), threading.TIMEOUT_MAX), + ) + drained = not ( + self._in_flight_op_count or self._unprovable_op_count or self._workers_alive + ) + if drained: + self._drained = True + left = self._sessions_left_open.items() + to_revoke = [key for key, by in left if by is _LeftOpen.FOR_DRAIN] if drained else [] + if to_revoke: + codes = self._metadata_call(self._store.batch_put_session_revoke, to_revoke) + with self._lock: + for key, code in zip(to_revoke, codes): + if _session_gone(code): + del self._sessions_left_open[key] + return drained + + def close(self) -> None: + """Closes the SDK client once a drain returned True and every region is unregistered. + Never called in production: it holds the GIL through unmount RPCs; process exit frees it.""" + with self._lock: + if self._closed: + return + if not self._drained: + raise DriverRejected("close needs a drain that returned True") + if any(r.live for r in self._registrations): + raise DriverRejected("close needs every registration unregistered") + self._closed = True + self._store.close() + + def stats(self) -> dict[str, int]: + """A snapshot for logs and tests; not part of the driver protocol.""" + with self._lock: + left = list(self._sessions_left_open.values()) + return { + **{ + f"rows_{r.status.value}_{reason}": self._skipped_rows[reason] + for reason, r in _SKIP_RESULTS.items() + }, + "ops_unprovable": self._unprovable_op_count, + "put_sessions_left_for_drain": left.count(_LeftOpen.FOR_DRAIN), + "put_sessions_left_to_master": left.count(_LeftOpen.TO_MASTER), + } + + # Admission, on the caller's thread. + + def _admit(self, rows: RowBatch, is_put: bool, on_change: OnChange | None) -> _Op: + """Checks ``rows`` and queues the operation, or raises ``DriverRejected`` having done + neither. The checks look at the batch once, whatever its size: its rows are laid out alike, + in the registration that holds their extent. The only work per row is the operation's list + of pending results, built at C speed.""" + if not len(rows): + raise DriverRejected("an operation needs at least one row") + if rows.spans_per_row != 1: + raise DriverRejected("this driver takes rows of one span") + if rows.row_bytes > MAX_ROW_BYTES: + raise DriverRejected(f"a row holds at most {MAX_ROW_BYTES} bytes") + per_part = max(1, min(self._rows_per_call, self._bytes_per_call // rows.row_bytes)) + with self._lock: + self._require_open(DriverRejected) + if self._draining: + raise DriverRejected("the driver is draining") + if not self._workers_alive: # every worker was stopped by an exception + raise DriverRejected("no worker thread is left to run the operation") + (extent,) = rows.extents + region = next( + (r for r in self._registrations if r.live and extent.within(r.span)), None + ) + if region is None: + raise DriverRejected(f"{extent} is outside every live registration") + op = _Op(self, is_put, rows, region, per_part, [_PENDING] * len(rows), on_change) + parts = -(-len(rows) // per_part) + if is_put: + if not self._queued_puts: # the wait GETS_BEFORE_PUT_IS_DUE measures starts now + self._gets_ahead_of_puts = 0 + self._queued_puts.append(op) + put_room = max(0, self._max_running_puts - self._running_puts) + self._work_due.notify(min(parts, put_room)) + else: + self._queued_gets.append(op) + self._work_due.notify(parts) + self._in_flight_op_count += 1 + region.ops_accessing += 1 + return op + + # The workers. + + def _part_is_takeable(self) -> bool: + """Under the lock: whether a worker could take a part now.""" + put_room = self._running_puts < self._max_running_puts + return bool(self._queued_gets) or (put_room and bool(self._queued_puts)) + + def _worker_loop(self) -> None: + """A worker thread: runs the parts ``_take`` gives it and waits while there is none; exits + once the driver drains. After each take it wakes an idle worker for a part still takeable. + The last worker to exit ends the rows no worker took, since none will.""" + try: + while True: + with self._lock: + part = self._take() + while part is None and not self._draining: + self._work_due.wait() + part = self._take() + if part is None: + return + if self._part_is_takeable(): + self._work_due.notify() + self._run_part(part) + finally: + done: list[_Op] = [] + with self._lock: + self._workers_alive -= 1 + # The last worker to exit ends the rows still queued, since no one will take + # them. After a drain nothing is queued; otherwise the others died of exceptions. + if not self._workers_alive: + done = self._end_untaken_rows(self._queued_ops(), _NO_WORKER_LEFT) + self._lock.notify_all() # drain waits for every worker thread to exit + for op in done: + self._call_on_change(op) + + def _take(self) -> _Part | None: + """Under the lock: the next part to run, or None when none may run now. Gets go first, but + once ``GETS_BEFORE_PUT_IS_DUE`` gets have gone ahead of a waiting put, that put is due and + goes next once a put may run; puts occupy at most ``_max_running_puts`` threads.""" + put_may_run = bool(self._queued_puts) and self._running_puts < self._max_running_puts + put_is_due = put_may_run and self._gets_ahead_of_puts >= GETS_BEFORE_PUT_IS_DUE + if self._queued_gets and not put_is_due: + if self._queued_puts: # this get goes ahead of a waiting put + self._gets_ahead_of_puts += 1 + return self._next_part(self._queued_gets) + if put_may_run: + self._gets_ahead_of_puts = 0 + self._running_puts += 1 + return self._next_part(self._queued_puts) + return None + + @staticmethod + def _next_part(queue: deque[_Op]) -> _Part: + """Under the lock: cuts the next part off the operation at the head of ``queue``; the + operation leaves the queue with its last rows.""" + op = queue[0] + start = op.first_untaken_row + op.first_untaken_row = min(start + op.rows_per_part, len(op.rows)) + op.parts_running += 1 + if op.all_taken: + queue.popleft() + return _Part(op, start, op.first_untaken_row) + + def _run_part(self, part: _Part) -> None: + """Runs ``part``'s SDK calls and records how they ended; calls the operation's + ``on_change`` once it is done or turned unprovable. A part whose calls stop before they + return its results still ends, its rows FAULT and its access unproven.""" + ended: _PartEnd | None = None + try: + ended = self._make_sdk_calls(part) + finally: + on_change_due = self._end_part(part, ended) + if on_change_due: + self._call_on_change(part.op) + + def _make_sdk_calls(self, part: _Part) -> _PartEnd: + direction = "put" if part.op.is_put else "get" + batch = part.op.rows[part.start : part.stop] + try: + rows = _Rows(batch.keys(), batch.addresses[:, 0].tolist(), batch.row_bytes) + except Exception as error: # no SDK call was made, so nothing touches the slots + failed = [faulted(f"{direction} failed before any SDK call: {error!r}")] * len(batch) + logger.warning(f"mooncake {direction}: {failed[0].fault}") + return _PartEnd(failed, (), access_ended=True) + started_s = time.monotonic() + try: + if part.op.is_put: + results, access_ended = self._put_part(rows) + else: + results, access_ended = self._get_part(rows) + except Exception as error: # a bug, not an SDK answer: a transfer may still be running + results, access_ended = _PartResults(), False + results.fault(rows.keys, f"{direction} failed: {error!r}") + if not results.all_ready or not access_ended: + try: + _warn(direction, rows, results, access_ended, time.monotonic() - started_s) + except Exception as error: # a log line must not change how the part ended + logger.warning(f"mooncake: warning about a part raised {error!r}") + skip_reasons = tuple(results.skip_reasons.values()) + return _PartEnd(results.in_order(rows.keys), skip_reasons, access_ended) + + def _end_part(self, part: _Part, ended: _PartEnd | None) -> bool: + """Records how ``part`` ended (``None``: its calls stopped before returning its results, + so its rows FAULT and its access is unproven); a part whose access did not end makes its + operation unprovable. True when the operation is done or turned unprovable here.""" + op = part.op + if ended is None: + stopped = faulted("the worker stopped before the row was decided") + ended = _PartEnd([stopped] * (part.stop - part.start), (), access_ended=False) + access_ended = ended.access_ended + with self._lock: + op.results[part.start : part.stop] = ended.results + self._skipped_rows.update(ended.skip_reasons) + if op.is_put: + self._running_puts -= 1 + if self._part_is_takeable(): # a waiting put may run now + self._work_due.notify() + turned_unprovable = not access_ended and not op.unprovable + if turned_unprovable: # holds its region for good + op.unprovable = True + self._unprovable_op_count += 1 + op.parts_running -= 1 + return self._end_op_if_done(op) or turned_unprovable + + def _queued_ops(self) -> list[_Op]: + """Under the lock: every operation with rows no worker took, gets first.""" + return [*self._queued_gets, *self._queued_puts] + + def _end_untaken_rows(self, ops: list[_Op], result: RowResult) -> list[_Op]: + """Under the lock: ends the rows of ``ops`` no worker took as ``result``, taking each + operation out of its queue; returns those this left done.""" + done = [] + for op in [op for op in ops if not op.all_taken]: + (self._queued_puts if op.is_put else self._queued_gets).remove(op) + untaken = len(op.rows) - op.first_untaken_row + op.results[op.first_untaken_row :] = [result] * untaken + op.first_untaken_row = len(op.rows) + if self._end_op_if_done(op): + done.append(op) + return done + + def _end_op_if_done(self, op: _Op) -> bool: + """Under the lock, after a part of ``op`` ended or its untaken rows did: once no row waits + and no part runs, the operation is done: every row concluded and its access ended, unless + unprovable. It then leaves ``_in_flight_op_count`` and, unless unprovable, releases its + region. Wakes ``drain``; True when the operation is done.""" + self._lock.notify_all() + if not op.all_taken or op.parts_running: + return False + self._in_flight_op_count -= 1 + if not op.unprovable: + op.region.ops_accessing -= 1 + return True + + def _call_on_change(self, op: _Op) -> None: + """Outside the lock, as ``OnChange`` asks.""" + if op.on_change is None: + return + try: + op.on_change(op) + except Exception as error: # it must not raise; a worker must live on if it does + logger.warning(f"mooncake: on_change raised {error!r}") + + # Put: start, one scatter, end what was written whole, revoke the rest. + + def _put_part(self, rows: _Rows) -> tuple[_PartResults, bool]: + """Returns the part's results and whether access ended. Whatever stops the calls, each + reserved key whose session they did not see closed is left open (class docstring, + step 3).""" + results = _PartResults() + reserved, held_by_left_session = self._reserve_keys(rows.keys, results) + # Reserved keys known to hold no session (never opened, or ended or revoked), and those + # whose object the target may still write into. + closed: set[str] = set() + to_master: set[str] = set() + try: + if held_by_left_session: # never put again, but another client may have written it + self._confirm_existing(held_by_left_session, results, "held_by_own_put") + if not reserved: + return results, True + access_ended = self._put_reserved(rows.only(reserved), results, closed, to_master) + return results, access_ended + finally: + self._release_keys_and_record_left_open(reserved, closed, to_master) + + def _reserve_keys(self, keys: list[str], results: _PartResults) -> tuple[list[str], list[str]]: + """Reserves every key that no put of this driver is writing and no session it left open + holds, and returns them, then the keys a left-open session holds. A key another put of + this driver is writing is NOT_READY.""" + with self._lock: + being_put, left_open = self._keys_being_put, self._sessions_left_open.keys() + if being_put.isdisjoint(keys) and left_open.isdisjoint(keys): # usually + reserved, held_by_left_session = keys, [] + else: + reserved, held_by_left_session = [], [] + for key in keys: + if key in being_put: + results.skip(key, "held_by_own_put") + elif key in left_open: + held_by_left_session.append(key) + else: + reserved.append(key) + self._keys_being_put.update(reserved) + return reserved, held_by_left_session + + def _release_keys_and_record_left_open( + self, reserved: list[str], closed: set[str], to_master: set[str] + ) -> None: + """Ends the reservation of ``reserved`` and records each session not ``closed`` as left + open: for the master when its key is in ``to_master``, otherwise for ``drain``.""" + left: dict[str, _LeftOpen] = {} + if not closed.issuperset(reserved): # usually every session ended + left = { + key: _LeftOpen.TO_MASTER if key in to_master else _LeftOpen.FOR_DRAIN + for key in reserved + if key not in closed + } + with self._lock: + self._keys_being_put.difference_update(reserved) + self._sessions_left_open.update(left) + if left: + first, why = next(iter(left.items())) + logger.warning( + f"mooncake: {len(left)} put sessions left open, first {first} {why.value}" + ) + + def _put_reserved( + self, rows: _Rows, results: _PartResults, closed: set[str], to_master: set[str] + ) -> bool: + """Starts a session per row, writes and finishes the started ones, and confirms the rows + whose object exists already. Returns whether access ended.""" + starts = self._metadata_call(self._store.batch_put_session_start, rows.keys, rows.sizes) + started, existing = rows.keys, [] + if any(starts): # not every session started (OK is 0) + started = [] + for key, code in zip(rows.keys, starts): + if code == _Code.OK: + started.append(key) + continue + if code != _Code.METADATA_CALL_FAILED: # the SDK answered: no session opened + closed.add(key) + if code == _Code.OBJECT_ALREADY_EXISTS: + existing.append(key) + elif code == _Code.NO_AVAILABLE_HANDLE: + results.skip(key, "pool_full") + else: + # Any other code faults the row. That includes -600: this driver holds no + # session of the key, so it is something else, such as a cachelib master's + # object cap (16 MiB - 16 B) or a replicate config the session path refuses + # (the driver passes none); the master's log says which. + results.fault([key], f"put session start: {_describe(code)}") + access_ended = True + if started: + access_ended = self._write_and_finish(rows.only(started), results, closed, to_master) + if existing: # readable only once the master shows it committed + self._confirm_existing(existing, results, "uncommitted") + return access_ended + + def _write_and_finish( + self, rows: _Rows, results: _PartResults, closed: set[str], to_master: set[str] + ) -> bool: + """Writes the rows whose put session started in one scatter, then ends the session of each + row written whole and revokes the rest, except those the target may still write into. + Returns whether access ended.""" + count = len(rows.keys) + try: + # One range per row: the whole payload, written at offset 0 of its object. + answers = _one_answer_per_key( + self._store.batch_put_from_multi_buffer_ranges( + rows.keys, + [[a] for a in rows.addresses], + [[rows.row_bytes]] * count, + [[0]] * count, + ), + count, + ) + except Exception as error: # the scatter may still be reading the slots + # Neither end (the object may be torn) nor revoke (the master would give its space + # to another object while the scatter may still write into it). + results.fault(rows.keys, f"batch_put_from_multi_buffer_ranges raised {error!r}") + to_master.update(rows.keys) + return False + written, to_revoke = rows.keys, [] + if answers != rows.sizes: # rare: some row was not written whole + written = [] + for key, answer in zip(rows.keys, answers): + if answer == rows.row_bytes: + written.append(key) + continue + # Its object must never be marked complete. After a TRANSFER_FAIL a TCP target + # may still write into the object's space (module docstring). + results.fault([key], f"put ranges: {_describe_moved(answer, rows.row_bytes)}") + if answer == _Code.TRANSFER_FAIL: + to_master.add(key) + else: + to_revoke.append(key) + if written: + to_revoke += self._end_sessions(written, results, closed) + if to_revoke: + codes = self._metadata_call(self._store.batch_put_session_revoke, to_revoke) + closed.update(key for key, code in zip(to_revoke, codes) if _session_gone(code)) + return True + + def _end_sessions(self, keys: list[str], results: _PartResults, closed: set[str]) -> list[str]: + """Ends the sessions of rows written whole. Returns the keys whose end was refused, to + revoke; a session whose end may still reach the master, answered ``_END_OUTCOME_UNKNOWN`` + first, stays open even when the retry was refused.""" + end = self._store.batch_put_session_end + codes = self._metadata_call(end, keys) + if not any(codes): # every session ended (OK is 0) + closed.update(keys) + return [] + final = dict(zip(keys, codes)) + retry = [key for key, code in final.items() if code in _END_RETRIED] + if retry: + final.update(zip(retry, self._metadata_call(end, retry))) + refused = [] + for key, first in zip(keys, codes): + code = final[key] + if code == _Code.OK: + closed.add(key) + continue + results.fault([key], f"put session end: {_describe(code)}") + if first not in _END_OUTCOME_UNKNOWN: + refused.append(key) + return refused + + def _confirm_existing( + self, keys: list[str], results: _PartResults, reason_if_absent: str + ) -> None: + """Leaves READY the rows whose objects the master shows complete; skips the others for + ``reason_if_absent``, or faults them when the master could not tell.""" + for key, answer in zip(keys, self._metadata_call(self._store.batch_is_exist, keys)): + if answer == OBJECT_ABSENT: + results.skip(key, reason_if_absent) + elif answer != OBJECT_EXISTS: + results.fault([key], f"existence check: {_describe(answer)}") + + # Get: one batch_get_into; rows whose read lease expired mid-read are read once more. + + def _get_part(self, rows: _Rows) -> tuple[_PartResults, bool]: + results = _PartResults() + return results, self._get_into(rows, results, is_reread=False) + + def _get_into(self, rows: _Rows, results: _PartResults, is_reread: bool) -> bool: + """Reads ``rows`` and records each one not read whole; returns whether access ended.""" + try: + answers = _one_answer_per_key( + self._store.batch_get_into(rows.keys, rows.addresses, rows.sizes), len(rows.keys) + ) + except Exception as error: # the transfers may still be writing the slots + results.fault(rows.keys, f"batch_get_into raised {error!r}") + return False + if answers == rows.sizes: # usually every row was read whole + return True + expired = [] + for key, answer in zip(rows.keys, answers): + if answer == rows.row_bytes: + continue + if answer == _Code.LEASE_EXPIRED and not is_reread: + expired.append(key) + elif answer == _Code.REPLICA_IS_NOT_READY and not is_reread: + results.skip(key, "no_complete_replica") + else: + results.recorded[key] = _result_of_get(answer, rows.row_bytes, is_reread) + # A TRANSFER_FAIL future may give up while its RDMA still runs (module docstring). + access_ended = _Code.TRANSFER_FAIL not in answers + if expired: + # The read lease expired before the transfer finished, so the master may have reused + # the object's space mid-read: read once more. The reread is READY only if it reads + # the whole object again; anything else, a second expiry included, is a FAULT. A + # TRANSFER_FAIL of the first read keeps access open whatever the reread does. + reread_access_ended = self._get_into(rows.only(expired), results, is_reread=True) + access_ended = access_ended and reread_access_ended + return access_ended + + # Helpers. + + def _metadata_call(self, call: Callable, keys: list[str], *args) -> list[int]: + """A session or existence call. It touches no row memory, so a raise or a wrong answer + count fails only its keys.""" + try: + return _one_answer_per_key(call(keys, *args), len(keys)) + except Exception as error: + logger.warning(f"mooncake: {getattr(call, '__name__', call)} failed: {error!r}") + return [_Code.METADATA_CALL_FAILED] * len(keys) + + def _add(self, registration: _Registration) -> None: + with self._lock: + self._registrations.append(registration) + + def _checked_registration(self, handle: _Registration) -> _Registration: + self._require_open() + if not isinstance(handle, _Registration) or handle not in self._registrations: + raise ValueError("handle does not belong to this driver") + return handle + + def _checked_op(self, op: _Op) -> _Op: + self._require_open() + if not isinstance(op, _Op) or op.owner is not self: + raise ValueError("operation does not belong to this driver") + return op + + def _require_open(self, error: type[Exception] = RuntimeError) -> None: + """Raises ``error`` once closed: DriverRejected refuses new work, RuntimeError the rest.""" + if self._closed: + raise error("the driver is closed") + + +def _result_of_get(answer: int, size: int, is_reread: bool) -> RowResult: + """The result of a row ``batch_get_into`` did not read whole.""" + if answer >= 0: # a read answers the object's size, which must be the row's + return faulted(f"object is {answer} bytes, expected {size}") + if is_reread and answer in (_Code.OBJECT_NOT_FOUND, _Code.REPLICA_IS_NOT_READY): + # The first read already wrote this slot, and MISS or NOT_READY would claim that nothing + # landed there. + return faulted(f"reread after LEASE_EXPIRED: {_describe(answer)}") + if answer == _Code.OBJECT_NOT_FOUND: + return _MISS + return faulted(f"batch_get_into: {_describe(answer)}") + + +def _session_gone(revoke_answer: int) -> bool: + """Whether a revocation left no session: revoked, or none was left (INVALID_PARAMS).""" + return revoke_answer in (_Code.OK, _Code.INVALID_PARAMS) + + +def _describe_moved(answer: int, size: int) -> str: + return f"moved {answer} of {size} bytes" if answer >= 0 else _describe(answer) + + +def _one_answer_per_key(answers: Sequence[int], count: int) -> list[int]: + answers = list(answers) + if len(answers) != count: + raise _WrongAnswerCount(f"{len(answers)} answers for {count} keys") + return answers + + +def _describe(code: int) -> str: + try: + return f"{_Code(code).name}({code})" + except ValueError: + return f"code({code})" + + +def _warn( + direction: str, rows: _Rows, results: _PartResults, access_ended: bool, seconds: float +) -> None: + """One warning per part that failed or skipped a row, or left access open; a MISS is + neither.""" + went_wrong = [(key, result) for key, result in results.recorded.items() if result is not _MISS] + if access_ended and not went_wrong: + return + first = "" + if went_wrong: + key, result = went_wrong[0] + why = result.fault or f"{result.status.value} ({results.skip_reasons[key]})" + first = f"; first {key}: {why}" + access = "ended" if access_ended else "may still be in progress" + logger.warning( + f"mooncake {direction}: {len(went_wrong)} of {len(rows.keys)} rows failed or skipped in " + f"{seconds:.3f} s{first}; access {access}" + ) diff --git a/tests/unittest/_torch/disaggregation/blob/fake_mooncake.py b/tests/unittest/_torch/disaggregation/blob/fake_mooncake.py new file mode 100644 index 000000000000..0bd9469f577a --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/fake_mooncake.py @@ -0,0 +1,299 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A fake ``mooncake.store.MooncakeDistributedStore`` that moves real bytes and can be scripted. + +It implements only the SDK methods the Mooncake driver may call (plus ``setup`` and ``close`` for +the opening code), so a call to any other one, a copying API for instance, fails with +``AttributeError``. Several fakes share one ``FakePool``, as clients of one master do. + +Unscripted calls follow SDK 0.3.13: + +- ``batch_put_session_start`` reserves an object on the pool: -600 when this client already holds + the key's session, -705 when the pool holds the key (committed or reserved). + ``batch_put_from_multi_buffer_ranges`` copies into the reservation (-600 without a session). + ``batch_put_session_end`` commits it and ``batch_put_session_revoke`` drops it; both answer -600 + without a session and keep the session when scripted to fail. +- ``batch_get_into`` answers each object's size after copying it (-704 absent, -703 reserved only, + -600 when the object is longer than its buffer). +- A buffer outside every registration fails with TRANSFER_FAIL, as RDMA does. + +``KeyGatedClient`` stops one method's calls naming a key at a gate, in front of any SDK client, +this fake or a real one. +""" + +from __future__ import annotations + +import ctypes +import threading +from collections import defaultdict, deque +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field + +OBJECT_EXISTS, OBJECT_ABSENT = 1, 0 +INTERNAL_ERROR = -1 +INVALID_PARAMS = -600 +REPLICA_IS_NOT_READY = -703 +OBJECT_NOT_FOUND = -704 +OBJECT_ALREADY_EXISTS = -705 +TRANSFER_FAIL = -800 +GATED_CALL_S = 50.0 # how long a blocked call waits for its gate: past any test's own wait + + +@dataclass +class FakePool: + """The objects of one Mooncake pool: committed ones, and reservations by put sessions.""" + + objects: dict[str, bytes] = field(default_factory=dict) + reserved: dict[str, bytearray] = field(default_factory=dict) + lock: threading.Lock = field(default_factory=threading.Lock) + + +@dataclass(frozen=True) +class SdkCall: + method: str + thread: int + args: tuple + kwargs: dict + answers: list = field(default_factory=list) # filled in once the call answers + + +class Gate: + """Holds one SDK call: ``entered`` is set when it arrives, it goes on once ``release`` is.""" + + def __init__(self) -> None: + self.entered = threading.Event() + self.release = threading.Event() + + +class KeyGatedClient: + """An SDK client whose ``method`` calls naming ``key`` wait inside the call until ``release`` is + set; every other call goes straight through.""" + + def __init__(self, store, method: str, key: str) -> None: + self.store = store + self.method = method + self.key = key + self.entered = threading.Event() + self.release = threading.Event() + + def __getattr__(self, name: str): + call = getattr(self.store, name) + if name != self.method: + return call + + def gated(keys, *args): + if self.key in keys: + self.entered.set() + # Past the suite's timeout: only the check's unblock lets it go. + self.release.wait(GATED_CALL_S) + return call(keys, *args) + + return gated + + +class FakeMooncakeStore: + """See the module docstring. ``script`` and ``block`` steer the next calls of a method.""" + + def __init__(self, pool: FakePool | None = None) -> None: + self.pool = pool if pool is not None else FakePool() + self.calls: list[SdkCall] = [] + self.registered: dict[int, int] = {} + self.put_sessions: set[str] = set() + self.setup_answer = 0 + self._lock = threading.Lock() + self._steps: dict[str, deque] = defaultdict(deque) + self._gates: dict[str, deque[Gate]] = defaultdict(deque) + + def script(self, method: str, step) -> None: + """Steers the next call of ``method`` that has no step yet. A list answers per key, + ``None`` meaning the real behaviour for that key; an exception is raised; a callable + rewrites the real answers (to drop or add one, say).""" + self._steps[method].append(step) + + def block(self, method: str) -> Gate: + """Holds the next call of ``method``, before it does anything, until the gate opens.""" + gate = Gate() + self._gates[method].append(gate) + return gate + + def calls_to(self, method: str) -> list[SdkCall]: + with self._lock: + return [call for call in self.calls if call.method == method] + + def keys_answered(self, method: str, answer: int) -> list[str]: + """Every key a call of ``method`` answered ``answer`` for, in call order.""" + with self._lock: + return [ + key + for call in self.calls + if call.method == method + for key, got in zip(call.args[0], call.answers) + if got == answer + ] + + # The SDK surface. + + def setup(self, *args, **kwargs) -> int: + self._record("setup", args, kwargs) + return self.setup_answer + + def close(self) -> int: + self._record("close", (), {}) + return 0 + + def register_buffer(self, buffer_ptr: int, size: int) -> int: + args = (buffer_ptr, size) + (answer,) = self._answer("register_buffer", args, {}, 1, lambda _: self._register(*args)) + return answer + + def unregister_buffer(self, buffer_ptr: int) -> int: + args = (buffer_ptr,) + (answer,) = self._answer( + "unregister_buffer", args, {}, 1, lambda _: self._unregister(*args) + ) + return answer + + def batch_put_session_start(self, *args, **kwargs) -> list[int]: + keys, sizes = args + return self._answer( + "batch_put_session_start", + args, + kwargs, + len(keys), + lambda i: self._start_put(keys[i], sizes[i]), + ) + + def batch_put_from_multi_buffer_ranges( + self, keys, all_buffer_ptrs, all_sizes, all_dst_offsets + ) -> list[int]: + args = (keys, all_buffer_ptrs, all_sizes, all_dst_offsets) + return self._answer( + "batch_put_from_multi_buffer_ranges", + args, + {}, + len(keys), + lambda i: self._write(keys[i], all_buffer_ptrs[i], all_sizes[i], all_dst_offsets[i]), + ) + + def batch_put_session_end(self, keys: Sequence[str]) -> list[int]: + return self._answer( + "batch_put_session_end", (keys,), {}, len(keys), lambda i: self._end_put(keys[i]) + ) + + def batch_put_session_revoke(self, keys: Sequence[str]) -> list[int]: + return self._answer( + "batch_put_session_revoke", (keys,), {}, len(keys), lambda i: self._revoke(keys[i]) + ) + + def batch_is_exist(self, keys: Sequence[str]) -> list[int]: + return self._answer( + "batch_is_exist", (keys,), {}, len(keys), lambda i: self._exists(keys[i]) + ) + + def batch_get_into(self, *args, **kwargs) -> list[int]: + keys, pointers, sizes = args + return self._answer( + "batch_get_into", + args, + kwargs, + len(keys), + lambda i: self._read(keys[i], pointers[i], sizes[i]), + ) + + # Real behaviour, one key at a time. + + def _register(self, pointer: int, size: int) -> int: + with self._lock: + self.registered[pointer] = size + return 0 + + def _unregister(self, pointer: int) -> int: + with self._lock: + return 0 if self.registered.pop(pointer, None) is not None else INTERNAL_ERROR + + def _start_put(self, key: str, size: int) -> int: + with self._lock, self.pool.lock: + if key in self.put_sessions: + return INVALID_PARAMS + if key in self.pool.objects or key in self.pool.reserved: + return OBJECT_ALREADY_EXISTS + self.pool.reserved[key] = bytearray(size) + self.put_sessions.add(key) + return 0 + + def _write(self, key: str, pointers: list[int], sizes: list[int], offsets: list[int]) -> int: + with self._lock, self.pool.lock: + target = self.pool.reserved.get(key) if key in self.put_sessions else None + if target is None or any(o + n > len(target) for o, n in zip(offsets, sizes)): + return INVALID_PARAMS + if not all(self._is_registered(p, n) for p, n in zip(pointers, sizes)): + return TRANSFER_FAIL + for pointer, size, offset in zip(pointers, sizes, offsets): + target[offset : offset + size] = ctypes.string_at(pointer, size) + return sum(sizes) + + def _end_put(self, key: str) -> int: + with self._lock, self.pool.lock: + if key not in self.put_sessions: + return INVALID_PARAMS + self.put_sessions.discard(key) + self.pool.objects[key] = bytes(self.pool.reserved.pop(key)) + return 0 + + def _revoke(self, key: str) -> int: + with self._lock, self.pool.lock: + if key not in self.put_sessions: + return INVALID_PARAMS + self.put_sessions.discard(key) + self.pool.reserved.pop(key, None) + return 0 + + def _exists(self, key: str) -> int: + with self.pool.lock: + return OBJECT_EXISTS if key in self.pool.objects else OBJECT_ABSENT + + def _read(self, key: str, pointer: int, size: int) -> int: + with self.pool.lock: + data, reserved = self.pool.objects.get(key), key in self.pool.reserved + if data is None: + return REPLICA_IS_NOT_READY if reserved else OBJECT_NOT_FOUND + if len(data) > size: + return INVALID_PARAMS + if not self._is_registered(pointer, size): + return TRANSFER_FAIL + ctypes.memmove(pointer, data, len(data)) + return len(data) + + def _is_registered(self, pointer: int, size: int) -> bool: + with self._lock: + return any( + start <= pointer and pointer + size <= start + length + for start, length in self.registered.items() + ) + + # Scripting. + + def _record(self, method: str, args: tuple, kwargs: dict) -> SdkCall: + call = SdkCall(method, threading.get_ident(), args, kwargs) + with self._lock: + self.calls.append(call) + return call + + def _answer(self, method, args, kwargs, count, real: Callable[[int], int]) -> list[int]: + call = self._record(method, args, kwargs) + with self._lock: + gate = self._gates[method].popleft() if self._gates[method] else None + step = self._steps[method].popleft() if self._steps[method] else None + if gate is not None: + gate.entered.set() + gate.release.wait(GATED_CALL_S) + if isinstance(step, BaseException): + raise step + if isinstance(step, list): + assert len(step) == count, f"a scripted {method} answers {count} keys" + answers = [real(i) if answer is None else answer for i, answer in enumerate(step)] + else: + answers = [real(i) for i in range(count)] + answers = step(answers) if callable(step) else answers + call.answers.extend(answers) + return answers diff --git a/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py new file mode 100644 index 000000000000..34bd94a5fe78 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py @@ -0,0 +1,1417 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Mooncake driver over a fake SDK that moves real bytes. + +Pins the mapping of SDK 0.3.13 answers to row results, the put session bookkeeping (every session +the driver starts is ended, revoked or left open; only rows written whole are ended), the access-end +rules (a put scatter that returns ends access, whatever it answers; a ``batch_get_into`` ends it +unless a row answers TRANSFER_FAIL; a call that raises or miscounts does not, and its access is +unprovable), registration, drain and close, the parts and their scheduling (parts cut by rows +and bytes, gets first with a waiting put after GETS_BEFORE_PUT_IS_DUE gets, puts leaving one +worker to gets, a fixed pool of workers, cancel and drain ending the rows no worker took, +``on_change`` once an operation is done), and a backend round trip over the fake SDK. The shared +``driver_conformance`` checks run over the fake SDK too. +""" + +from __future__ import annotations + +import math +import threading +import time +from collections import Counter +from collections.abc import Iterator +from typing import NamedTuple + +import numpy as np +import pytest +from cpu_fixtures import ( + PROMPT_S, + WAIT_S, + CpuRegion, + CpuView, + Endpoint, + FakeHold, + Rows, + WakeCounter, + object_key, + store_key, + wait_until, +) +from fake_mooncake import FakeMooncakeStore, FakePool, KeyGatedClient + +from tensorrt_llm._torch.disaggregation.base.backend import Attempt, Backend, Capabilities +from tensorrt_llm._torch.disaggregation.blob import mooncake as mooncake_module +from tensorrt_llm._torch.disaggregation.blob import store as store_module +from tensorrt_llm._torch.disaggregation.blob.backend import BlobBackend, BlobConfig, RegionBinding +from tensorrt_llm._torch.disaggregation.blob.mooncake import MooncakeDriver +from tensorrt_llm._torch.disaggregation.blob.store import ( + DriverCapabilities, + DriverRejected, + RegionLayout, + RegistrationError, + RowBatch, + RowIO, + RowStatus, + Span, + StoreDriver, + StoreError, + concluded, +) +from tensorrt_llm._torch.disaggregation.testing.driver_conformance import DriverConformance +from tensorrt_llm._torch.disaggregation.testing.host_memory import HostBuffer + +# The driver's workers run until it drains; ``driver_workers_return`` drains every driver a test +# left open and waits for its workers, instead of the thread-leak check. +# An exception that escapes a thread fails the test; a test that ends a thread on purpose +# requests the ``thread_escapes`` fixture. +pytestmark = [ + pytest.mark.cpu_only, + pytest.mark.threadleak(enabled=False), + pytest.mark.filterwarnings("error::pytest.PytestUnhandledThreadExceptionWarning"), +] + +SLOT = 8 +PUT_START = "batch_put_session_start" +PUT_RANGES = "batch_put_from_multi_buffer_ranges" +PUT_END = "batch_put_session_end" +PUT_REVOKE = "batch_put_session_revoke" +GET = "batch_get_into" +DATA_CALL = {"put": PUT_RANGES, "get": GET} +METADATA_CALL_FAILED = "existence check: METADATA_CALL_FAILED(-100000)" + + +@pytest.fixture(autouse=True) +def driver_workers_return(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + drivers = [] + original_init = MooncakeDriver.__init__ + + def tracked_init(self, *args, **kwargs) -> None: + drivers.append(self) + original_init(self, *args, **kwargs) + + monkeypatch.setattr(MooncakeDriver, "__init__", tracked_init) + yield + for driver in drivers: + try: + driver.drain(0.0) + except Exception: # closed already, or a test's scripted failure + pass + wait_until( + lambda: not any(t.name.startswith("mooncake-driver-") for t in threading.enumerate()), + "the driver's workers to return", + ) + + +def key(name: str) -> str: + return f"t/v1/{name}" + + +class Bench: + """A ``MooncakeDriver`` over a fake SDK, with ``slots`` registered host slots; ``tuning`` goes + to the driver.""" + + def __init__(self, slots: int = 4, pool: FakePool | None = None, **tuning: int) -> None: + self.fake = FakeMooncakeStore(pool) + self.driver = MooncakeDriver(self.fake, **tuning) + self.buffer = HostBuffer(SLOT * slots, fill=0x2D) + layout = RegionLayout(SLOT, slots) + self.handle = self.driver.register(self.buffer.address, self.buffer.nbytes, layout) + + def rows(self, *names: str) -> RowBatch: + return RowBatch.of( + [RowIO(key(n), (self.buffer.span(i * SLOT, SLOT),)) for i, n in enumerate(names)] + ) + + def store(self, **objects: bytes) -> None: + for name, data in objects.items(): + self.fake.pool.objects[key(name)] = data + + def write(self, slot: int, data: bytes) -> None: + self.buffer.write(data, slot * SLOT) + + def read(self, slot: int) -> bytes: + return self.buffer.read(slot * SLOT, SLOT) + + def run(self, direction: str, *names: str, on_change=None): + """The results and access end of a put or get of ``names`` once it settled; ``last`` is its + handle.""" + rows = self.rows(*names) + if direction == "put": + self.last = self.driver.put(rows, on_change=on_change) + else: + self.last = self.driver.get(rows, on_change=on_change) + return self.settle(self.last) + + def settle(self, op): + """``op``'s results and access end, once its results concluded and its access ended or + can never be proven to.""" + driver = self.driver + wait_until( + lambda: ( + concluded(driver.poll(op)) + and (driver.access_ended(op) or driver.access_unprovable(op)) + ), + "the operation to settle", + ) + return driver.poll(op), driver.access_ended(op) + + def fill_sources(self, direction: str, *names: str) -> None: + """Gives a put of ``names`` its slots' bytes, or a get its objects in the pool.""" + for slot, name in enumerate(names): + data = name.encode() * SLOT + self.write(slot, data) if direction == "put" else self.store(**{name: data}) + + def keys_of(self, method: str) -> list[str]: + return [k for call in self.fake.calls_to(method) for k in call.args[0]] + + +def summary(results) -> list[tuple[str, str | None]]: + return [(result.status.value, result.fault) for result in results] + + +def test_the_driver_is_a_store_driver_with_mooncake_capabilities() -> None: + bench = Bench() + assert isinstance(bench.driver, StoreDriver) + assert bench.driver.capabilities == DriverCapabilities( + get=True, + put=True, + cancel=True, + probe=False, + max_rows_per_op=None, + hints=False, + absent_is_miss=True, + device_memory=False, + multi_span=False, + ) + + +class DriverLeavingOperationsUnstarted(MooncakeDriver): + """A Mooncake driver whose workers never take a part.""" + + def _take(self): + return None + + +class TestMooncakeDriverConformance(DriverConformance): + """The shared driver checks over the fake SDK, with every scenario it can stage. Probe is not + declared, so the blocked and failing presence queries are not staged.""" + + def make_driver(self) -> StoreDriver: + # A part per row, so every operation of more than one row runs in several parts. + return MooncakeDriver(FakeMooncakeStore(), rows_per_call=1) + + def make_blocked_get(self): + sdk = KeyGatedClient(FakeMooncakeStore(), GET, key("blocked")) + return MooncakeDriver(sdk), key("blocked"), sdk.entered, sdk.release.set + + def make_blocked_put(self): + sdk = KeyGatedClient(FakeMooncakeStore(), PUT_RANGES, key("blocked")) + return MooncakeDriver(sdk), key("blocked"), sdk.entered, sdk.release.set + + def make_held_key(self): + pool = FakePool() + writer = FakeMooncakeStore(pool) # another client, still writing the key + assert writer.batch_put_session_start([key("held")], [SLOT]) == [0] + + def release() -> None: + writer.batch_put_session_revoke([key("held")]) + + return MooncakeDriver(FakeMooncakeStore(pool)), key("held"), release + + def make_full_store(self): + fake = FakeMooncakeStore() + # Room for one more object as the check uses it: its one two-row put finds room for the + # first only. + fake.script(PUT_START, [None, -200]) + return MooncakeDriver(fake) + + def make_crash_after_record(self): + fake = FakeMooncakeStore() + fake.script(PUT_START, RuntimeError("the SDK crashed")) + fake.script(GET, RuntimeError("the SDK crashed")) + return MooncakeDriver(fake) + + def make_driver_leaving_operations_unstarted(self): + return DriverLeavingOperationsUnstarted(FakeMooncakeStore()) + + def make_failing_get(self): + fake = FakeMooncakeStore() + fake.script(GET, [-900]) # the next read fails with RPC_FAIL; the check reads the key once + return MooncakeDriver(fake), key("failing") + + def make_unprovable_get(self): + fake = FakeMooncakeStore() + fake.script(GET, [-800]) # a transfer future gave up while its transfer may still run + return MooncakeDriver(fake), lambda op: None # the scripted answer concludes it + + +# Parts and their scheduling. + + +PART_LIMITS = { + # name: (rows_per_call, bytes_per_call, rows in each part) + "bytes": (10, 2 * SLOT + 1, [2] * 5), + "rows": (4, 10 * SLOT, [4, 4, 2]), + "whichever_comes_first": (3, 2 * SLOT, [2] * 5), + "a_row_over_bytes_per_call_goes_alone": (10, 1, [1] * 10), + "no_cut": (10, 10 * SLOT, [10]), +} + + +@pytest.mark.parametrize("case", PART_LIMITS) +def test_an_operation_of_any_size_runs_in_parts_cut_by_rows_and_bytes(case: str) -> None: + rows_per_call, bytes_per_call, expected = PART_LIMITS[case] + bench = Bench(slots=10, workers=1, rows_per_call=rows_per_call, bytes_per_call=bytes_per_call) + names = list("ABCDEFGHIJ") # one byte each: fill_sources repeats it over the slot + bench.fill_sources("put", *names) + results, access_ended = bench.run("put", *names) + assert summary(results) == [("ready", None)] * 10 and access_ended + assert [len(call.args[0]) for call in bench.fake.calls_to(PUT_RANGES)] == expected + assert all(bench.fake.pool.objects[key(name)] == name.encode() * SLOT for name in names) + + +def test_the_driver_refuses_tuning_that_is_not_a_positive_int() -> None: + for name in ("workers", "rows_per_call", "bytes_per_call"): + for value in (0, True, 1.5): + with pytest.raises(ValueError, match=f"{name} must be a positive int"): + MooncakeDriver(FakeMooncakeStore(), **{name: value}) + + +def data_calls(bench: Bench) -> list[str]: + """``get`` or ``put`` for each SDK call that moves rows or starts to, in call order: a + ``batch_get_into``, or the session start of a put.""" + names = {GET: "get", PUT_START: "put"} + return [names[call.method] for call in bench.fake.calls if call.method in names] + + +def test_a_waiting_put_goes_next_once_it_is_due() -> None: + bench = Bench(slots=3 * mooncake_module.GETS_BEFORE_PUT_IS_DUE + 3, workers=1) + fetched = [f"F{i}" for i in range(3 * mooncake_module.GETS_BEFORE_PUT_IS_DUE + 1)] + bench.store(**{name: bytes(SLOT) for name in fetched}) + gate = bench.fake.block(GET) + ops = [bench.driver.get(bench.rows(fetched[0]))] + assert gate.entered.wait(WAIT_S) # the only worker is inside the SDK + # Each row is an operation of its own; the slots are reused, as no assertion reads them. + ops += [bench.driver.get(bench.rows(name)) for name in fetched[1:]] + ops += [bench.driver.put(bench.rows(name)) for name in ("P1", "P2")] + gate.release.set() + for op in ops: + bench.settle(op) + # The blocked get was taken before any put waited, so it does not count toward one's turn. + gets_then_a_put = ["get"] * mooncake_module.GETS_BEFORE_PUT_IS_DUE + ["put"] + expected = ["get", *gets_then_a_put * 2] + assert data_calls(bench)[: len(expected)] == expected + + +def test_only_gets_taken_while_a_put_waits_count_toward_its_turn() -> None: + bench = Bench(slots=2, workers=1) + bench.store(**{f"F{i}": bytes(SLOT) for i in range(mooncake_module.GETS_BEFORE_PUT_IS_DUE + 2)}) + for i in range(mooncake_module.GETS_BEFORE_PUT_IS_DUE): # taken with no put waiting + bench.run("get", f"F{i}") + gate = bench.fake.block(GET) + held = bench.driver.get(bench.rows(f"F{mooncake_module.GETS_BEFORE_PUT_IS_DUE}")) + assert gate.entered.wait(WAIT_S) + put = bench.driver.put(bench.rows("P")) + late = bench.driver.get(bench.rows(f"F{mooncake_module.GETS_BEFORE_PUT_IS_DUE + 1}")) + gate.release.set() + for op in (held, put, late): + bench.settle(op) + gets = ["get"] * (mooncake_module.GETS_BEFORE_PUT_IS_DUE + 2) + assert data_calls(bench) == [*gets, "put"] # the late get goes first + + +def test_a_cancelled_put_leaves_its_turn_to_no_later_put() -> None: + due = mooncake_module.GETS_BEFORE_PUT_IS_DUE + bench = Bench(slots=4, workers=1) + fetched = [f"F{i}" for i in range(due + 2)] + bench.store(**{name: bytes(SLOT) for name in fetched}) + bench.fill_sources("put", "P", "Q") + gates = [bench.fake.block(GET) for _ in range(due + 1)] + ops = [bench.driver.get(bench.rows(fetched[0]))] + assert gates[0].entered.wait(WAIT_S) # taken before any put waited + put = bench.driver.put(bench.rows("P")) + ops += [bench.driver.get(bench.rows(name)) for name in fetched[1 : due + 1]] + for gate in gates[:due]: + gate.release.set() + assert gates[due].entered.wait(WAIT_S) # the gets ahead of P made it due + bench.driver.cancel(put) + calls_before = len(data_calls(bench)) + rows = bench.rows("Q", fetched[due + 1]) + ops += [bench.driver.put(rows[:1]), bench.driver.get(rows[1:])] + gates[due].release.set() + for op in ops: + bench.settle(op) + # No get went ahead of Q, which started waiting with no other put queued: gets go first. + assert data_calls(bench)[calls_before:] == ["get", "put"] + + +def test_puts_leave_one_worker_to_gets() -> None: + bench = Bench(slots=3, workers=2) + bench.fill_sources("put", "P1", "P2") + bench.store(F=b"ffffffff") + gate = bench.fake.block(PUT_RANGES) + rows = bench.rows("P1", "P2", "F") + p1, p2 = bench.driver.put(rows[:1]), bench.driver.put(rows[1:2]) + try: + assert gate.entered.wait(WAIT_S) + time.sleep(PROMPT_S) # P2 would have taken the other worker by now + assert bench.keys_of(PUT_START) == [key("P1")] + f = bench.driver.get(rows[2:]) + assert summary(bench.settle(f)[0]) == [("ready", None)] + time.sleep(PROMPT_S) # the worker F freed would have taken P2 by now + assert bench.keys_of(PUT_START) == [key("P1")] + finally: + gate.release.set() + for op in (p1, p2): + assert summary(bench.settle(op)[0]) == [("ready", None)] + + +class OnChangeRecorder: + """An ``on_change`` that keeps the handle it is called with and what ``poll`` then shows.""" + + def __init__(self, driver: MooncakeDriver) -> None: + self.driver = driver + self.handles: list = [] + self.shown: list[list[tuple]] = [] + + def __call__(self, op) -> None: + self.handles.append(op) + self.shown.append(summary(self.driver.poll(op))) + + def shown_for(self, op) -> list[list[tuple]]: + """What ``poll`` showed at each call, each of which named ``op``.""" + assert all(handle is op for handle in self.handles) + return self.shown + + +def test_a_cancel_ends_only_the_rows_no_worker_took() -> None: + bench = Bench(workers=1, bytes_per_call=SLOT) + bench.store(A=b"aaaaaaaa", B=b"bbbbbbbb", C=b"cccccccc") + gate = bench.fake.block(GET) + recorder = OnChangeRecorder(bench.driver) + op = bench.driver.get(bench.rows("A", "B", "C"), on_change=recorder) + try: + assert gate.entered.wait(WAIT_S) + bench.driver.cancel(op) + shown = summary(bench.driver.poll(op)) + assert shown == [("pending", None), ("cancelled", None), ("cancelled", None)] + assert not bench.driver.access_ended(op) # A's part is still in the SDK + assert recorder.shown == [] # A has not concluded, so there is nothing to act on yet + finally: + gate.release.set() + results, access_ended = bench.settle(op) + assert summary(results) == [("ready", None), ("cancelled", None), ("cancelled", None)] + assert access_ended and bench.keys_of(GET) == [key("A")] + wait_until(lambda: recorder.shown, "on_change once A's part ended") + assert recorder.shown_for(op) == [summary(results)] + assert bench.read(1) == bench.read(2) == b"-" * SLOT + bench.driver.cancel(op) # nothing left to cancel + assert bench.driver.drain(WAIT_S) is True + + +def test_drain_cancels_the_rows_no_worker_took() -> None: + bench = Bench(workers=1, bytes_per_call=SLOT) + bench.store(A=b"aaaaaaaa", B=b"bbbbbbbb") + gate = bench.fake.block(GET) + op = bench.driver.get(bench.rows("A", "B")) + try: + assert gate.entered.wait(WAIT_S) + assert bench.driver.drain(0.0) is False # A's part is still in the SDK + finally: + gate.release.set() + results, access_ended = bench.settle(op) + assert summary(results) == [("ready", None), ("cancelled", None)] and access_ended + assert bench.keys_of(GET) == [key("A")] + assert bench.driver.drain(WAIT_S) is True + + +def test_a_part_whose_access_never_ends_leaves_the_rest_of_its_operation_to_run() -> None: + bench = Bench(workers=1, bytes_per_call=SLOT) # a part per row + bench.store(A=b"aaaaaaaa", B=b"bbbbbbbb") + bench.fake.script(GET, [-800]) # A's transfer may still run + results, access_ended = bench.run("get", "A", "B") + assert [status for status, _ in summary(results)] == ["fault", "ready"] + assert not access_ended and bench.driver.access_unprovable(bench.last) + assert bench.keys_of(GET) == [key("A"), key("B")] + + +def test_an_on_change_that_raises_leaves_the_driver_serving() -> None: + bench = Bench(workers=1) + bench.store(A=b"aaaaaaaa") + raised = [] + + def raise_once(*_) -> None: + if not raised: + raised.append(True) + raise KeyError("boom") + + for _ in range(2): # the second runs only if the first left the worker serving + results, access_ended = bench.run("get", "A", on_change=raise_once) + assert summary(results) == [("ready", None)] and access_ended + assert raised and bench.driver.drain(WAIT_S) is True + + +def test_a_warning_that_raises_leaves_the_part_as_it_ended(monkeypatch: pytest.MonkeyPatch) -> None: + bench = Bench(workers=1) + + def broken(*_) -> None: + raise KeyError("a logging bug") + + monkeypatch.setattr(mooncake_module, "_warn", broken) + bench.fake.script(GET, [-703]) # NOT_READY: a skipped row, which the part warns of + results, access_ended = bench.run("get", "A") + assert summary(results) == [("not_ready", None)] and access_ended + assert bench.driver.drain(WAIT_S) is True + + +def test_the_driver_calls_on_change_once_the_operation_is_done() -> None: + bench = Bench(workers=1, bytes_per_call=SLOT) # a part per row, run one after the other + bench.store(A=b"aaaaaaaa", B=b"bbbbbbbb") + recorder = OnChangeRecorder(bench.driver) + op = bench.driver.get(bench.rows("A", "B"), on_change=recorder) + bench.settle(op) + wait_until(lambda: recorder.shown, "on_change once B's part ended") + # Not after A's part: a caller does nothing with some rows concluded. + assert recorder.shown_for(op) == [[("ready", None), ("ready", None)]] + + +def get_into_slot(bench: Bench, name: str, slot: int) -> RowBatch: + """One get row of ``name`` read into ``slot``.""" + return RowBatch.of([RowIO(key(name), (bench.buffer.span(slot * SLOT, SLOT),))]) + + +def test_a_get_goes_ahead_of_a_due_put_that_has_no_worker_to_run_on() -> None: + bench = Bench(slots=3, workers=2) # puts occupy at most one worker + bench.fill_sources("put", "P1", "P2") + names = [f"F{i}" for i in range(mooncake_module.GETS_BEFORE_PUT_IS_DUE + 2)] + bench.store(**{name: bytes(SLOT) for name in names}) + gate = bench.fake.block(PUT_RANGES) + rows = bench.rows("P1", "P2") + p1, p2 = bench.driver.put(rows[:1]), bench.driver.put(rows[1:]) + try: + assert gate.entered.wait(WAIT_S) # P1 holds the put worker; P2 waits for it + # Past GETS_BEFORE_PUT_IS_DUE gets P2 is due, but it still cannot run: the gets go on. + for name in names: + op = bench.driver.get(get_into_slot(bench, name, 2)) + assert summary(bench.settle(op)[0]) == [("ready", None)] + finally: + gate.release.set() + for op in (p1, p2): + assert summary(bench.settle(op)[0]) == [("ready", None)] + + +def test_gets_that_go_ahead_of_a_put_waiting_for_a_put_worker_count_toward_its_turn() -> None: + n = mooncake_module.GETS_BEFORE_PUT_IS_DUE + bench = Bench(slots=3, workers=2) # puts occupy at most one worker + bench.fill_sources("put", "P1", "P2") + names = [f"F{i}" for i in range(n + 2)] + bench.store(**{name: bytes(SLOT) for name in names}) + put_gate, get_gate = bench.fake.block(PUT_RANGES), None + rows = bench.rows("P1", "P2") + p1, p2 = bench.driver.put(rows[:1]), bench.driver.put(rows[1:]) + try: + assert put_gate.entered.wait(WAIT_S) # P1 holds the put worker; P2 waits for it + for name in names[:n]: # each goes ahead of P2, on the other worker + bench.settle(bench.driver.get(get_into_slot(bench, name, 2))) + get_gate = bench.fake.block(GET) + held = bench.driver.get(get_into_slot(bench, names[n], 2)) + assert get_gate.entered.wait(WAIT_S) # the other worker is busy too + late = bench.driver.get(get_into_slot(bench, names[n + 1], 2)) + put_gate.release.set() # P1's worker is the only one free: P2's turn came, so P2 first + for op in (p1, p2, late): + bench.settle(op) + assert data_calls(bench)[-2:] == ["put", "get"] + finally: + put_gate.release.set() + if get_gate is not None: + get_gate.release.set() + bench.settle(held) + + +def test_worker_threads_never_keep_the_process_alive() -> None: + bench = Bench() + workers = [t for t in threading.enumerate() if t.name.startswith("mooncake-driver-")] + assert workers and all(thread.daemon for thread in workers) + assert bench.driver.drain(WAIT_S) is True + + +def test_drain_waits_for_every_worker_thread_to_exit() -> None: + bench = Bench(workers=1) + bench.store(A=b"aaaaaaaa") + in_on_change, leave_on_change = threading.Event(), threading.Event() + + def held_on_change(_op) -> None: + in_on_change.set() + leave_on_change.wait(WAIT_S) + + op = bench.driver.get(bench.rows("A"), on_change=held_on_change) + assert in_on_change.wait(WAIT_S) # the part ended; its worker thread is inside on_change + try: + assert bench.driver.access_ended(op) + assert bench.driver.drain(0.05) is False # the worker thread has not exited + threading.Timer(0.1, leave_on_change.set).start() + started = time.monotonic() + assert bench.driver.drain(WAIT_S) is True + assert time.monotonic() - started < WAIT_S / 2 # woken as the thread exited + finally: + leave_on_change.set() + + +def test_the_conformance_suite_runs_its_checks_of_a_get_held_inside_the_sdk() -> None: + # put and get return before the SDK answers, so those checks must run rather than skip. + assert TestMooncakeDriverConformance.returns_before_the_store_answers + + +def test_drain_calls_on_change_for_the_operations_it_ended_before_it_waits() -> None: + bench = Bench(workers=1) + bench.store(A=b"aaaaaaaa", B=b"bbbbbbbb") + gate = bench.fake.block(GET) + rows = bench.rows("A", "B") + running = bench.driver.get(rows[:1]) + assert gate.entered.wait(WAIT_S) + recorder = OnChangeRecorder(bench.driver) + queued = bench.driver.get(rows[1:], on_change=recorder) + drained: list[bool] = [] + drainer = threading.Thread(target=lambda: drained.append(bench.driver.drain(WAIT_S))) + drainer.start() + try: + wait_until(lambda: recorder.shown, "on_change once B was cancelled") + assert recorder.shown_for(queued) == [[("cancelled", None)]] + assert not drained # still waiting for A's part, which is in the SDK + finally: + gate.release.set() + drainer.join(WAIT_S) + assert drained == [True] + assert summary(bench.settle(running)[0]) == [("ready", None)] + + +def test_drain_waits_out_a_budget_of_any_size() -> None: + bench = Bench(workers=1) + bench.store(A=b"aaaaaaaa") + gate = bench.fake.block(GET) + op = bench.driver.get(bench.rows("A")) + assert gate.entered.wait(WAIT_S) + threading.Timer(PROMPT_S, gate.release.set).start() + assert bench.driver.drain(math.inf) is True # once A's part ended + assert summary(bench.settle(op)[0]) == [("ready", None)] + + +@pytest.fixture +def thread_escapes() -> Iterator[list[BaseException]]: + """Catches the exceptions that end threads while a test that ends one on purpose runs, and + checks afterwards that each was a ``SystemExit``. The hook is put back before the drivers are + drained, so an escape there still reaches pytest and fails the test.""" + escapes: list[BaseException] = [] + pytest_hook = threading.excepthook + threading.excepthook = lambda args: escapes.append(args.exc_value) + try: + yield escapes + finally: + threading.excepthook = pytest_hook + assert all(isinstance(escape, SystemExit) for escape in escapes), escapes + + +def wait_for_an_escape(escapes: list[BaseException]) -> None: + """Waits until a thread ended by an exception: a thread runs its excepthook only after the + driver saw it leave, so a test that ends a thread waits for this before it ends.""" + wait_until(lambda: escapes, "a thread to end by an exception") + + +def test_a_part_whose_worker_thread_stops_faults_and_the_other_workers_go_on( + thread_escapes: list[BaseException], +) -> None: + bench = Bench(workers=2, bytes_per_call=SLOT) # a part per row + bench.store(A=b"aaaaaaaa", B=b"bbbbbbbb") + bench.fake.script(GET, SystemExit()) # escapes every handler: A's worker thread ends + results, access_ended = bench.run("get", "A", "B") + assert summary(results) == [ + ("fault", "the worker stopped before the row was decided"), + ("ready", None), # on the other worker + ] + assert not access_ended and bench.driver.access_unprovable(bench.last) + wait_for_an_escape(thread_escapes) + + +def test_once_no_worker_is_left_queued_rows_fault_and_operations_are_rejected( + thread_escapes: list[BaseException], +) -> None: + bench = Bench(workers=1, bytes_per_call=SLOT) # a part per row + bench.store(A=b"aaaaaaaa", B=b"bbbbbbbb") + bench.fake.script(GET, SystemExit()) # A's part ends the only worker thread + recorder = OnChangeRecorder(bench.driver) + op = bench.driver.get(bench.rows("A", "B"), on_change=recorder) + wait_until(lambda: len(recorder.shown) == 2, "A's end, then B's") + assert recorder.shown_for(op)[-1] == [ + ("fault", "the worker stopped before the row was decided"), + ("fault", "no worker thread is left to run the row"), # B waited for a worker in vain + ] + assert bench.driver.access_unprovable(op) # A's SDK call may still touch its slot + with pytest.raises(DriverRejected, match="no worker thread is left"): + bench.driver.get(bench.rows("A")) + wait_for_an_escape(thread_escapes) + + +def test_a_worker_that_cannot_start_stops_the_ones_already_started( + monkeypatch: pytest.MonkeyPatch, +) -> None: + start = threading.Thread.start + started: list[threading.Thread] = [] + + def refuse_the_second(thread: threading.Thread) -> None: + if thread.name.startswith("mooncake-driver-"): + if started: + raise RuntimeError("can't start new thread") + started.append(thread) + start(thread) + + monkeypatch.setattr(threading.Thread, "start", refuse_the_second) + with pytest.raises(RuntimeError, match="can't start new thread"): + MooncakeDriver(FakeMooncakeStore(), workers=2) + (first,) = started + first.join(WAIT_S) + assert not first.is_alive() + + +def test_drain_wakes_as_a_worker_thread_leaves_on_an_escape( + thread_escapes: list[BaseException], +) -> None: + bench = Bench(workers=1) + bench.store(A=b"aaaaaaaa") + leave = threading.Event() + + def tell_then_escape(*_) -> None: + leave.wait(WAIT_S) + raise SystemExit # past every handler: the thread exits after its part ended + + results, access_ended = bench.run("get", "A", on_change=tell_then_escape) + assert summary(results) == [("ready", None)] and access_ended + threading.Timer(0.1, leave.set).start() + started = time.monotonic() + assert bench.driver.drain(WAIT_S) is True + assert time.monotonic() - started < WAIT_S / 2 # woken as the thread left, not at the budget + wait_for_an_escape(thread_escapes) + + +# Answers to row results. + + +def skipped_counts(bench: Bench) -> dict[str, int]: + return {name: n for name, n in bench.driver.stats().items() if name.startswith("rows_") and n} + + +GET_ANSWERS = { + # name: (batch_get_into's answer (None: real), status, fault, the stat counting a skipped row) + "object_size": (None, "ready", None, None), # unscripted + "short_object": (5, "fault", "object is 5 bytes, expected 8", None), + "OBJECT_NOT_FOUND": (-704, "miss", None, None), + "REPLICA_IS_NOT_READY": (-703, "not_ready", None, "rows_not_ready_no_complete_replica"), + "INVALID_REPLICA": (-702, "fault", "batch_get_into: INVALID_REPLICA(-702)", None), + "INVALID_PARAMS": (-600, "fault", "batch_get_into: INVALID_PARAMS(-600)", None), + "RPC_FAIL": (-900, "fault", "batch_get_into: RPC_FAIL(-900)", None), + "TRANSFER_FAIL": (-800, "fault", "batch_get_into: TRANSFER_FAIL(-800)", None), + "unknown_code": (-9999, "fault", "batch_get_into: code(-9999)", None), +} + + +@pytest.mark.parametrize("case", GET_ANSWERS) +def test_get_answers_map_to_row_statuses(case: str) -> None: + answer, status, fault, counted = GET_ANSWERS[case] + bench = Bench() + bench.store(A=b"aaaaaaaa") + bench.fake.script(GET, [answer]) + results, access_ended = bench.run("get", "A") + assert summary(results) == [(status, fault)] + assert skipped_counts(bench) == ({counted: 1} if counted else {}) + assert (bench.read(0) == b"aaaaaaaa") == (status == "ready") + assert access_ended is (answer != -800) + assert bench.driver.access_unprovable(bench.last) is (answer == -800) + + +PUT_START_ANSWERS = { + # name: (the session start's answer, status, fault, the stat counting a skipped row) + "NO_AVAILABLE_HANDLE": (-200, "declined", None, "rows_declined_pool_full"), + # This driver holds no session of the key: not a session of its own, but, say, a cachelib + # master's object cap. + "INVALID_PARAMS": (-600, "fault", "put session start: INVALID_PARAMS(-600)", None), + "INVALID_REPLICA": (-702, "fault", "put session start: INVALID_REPLICA(-702)", None), + "RPC_FAIL": (-900, "fault", "put session start: RPC_FAIL(-900)", None), + "unknown_code": (-9999, "fault", "put session start: code(-9999)", None), +} + + +@pytest.mark.parametrize("case", PUT_START_ANSWERS) +def test_put_session_start_answers_map_to_row_statuses(case: str) -> None: + answer, status, fault, counted = PUT_START_ANSWERS[case] + bench = Bench() + bench.fake.script(PUT_START, [answer]) + results, access_ended = bench.run("put", "A") + assert summary(results) == [(status, fault)] + assert skipped_counts(bench) == ({counted: 1} if counted else {}) + assert access_ended + (start,) = bench.fake.calls_to(PUT_START) + assert len(start.args) == 2 and start.kwargs == {} # never a replicate config + for method in (PUT_RANGES, PUT_END, PUT_REVOKE, "batch_is_exist"): + assert bench.fake.calls_to(method) == [] # a session it did not start is not its own + + +PUT_WRITE_ANSWERS = { + # case: (the scatter's answer, the end's answer (None: real), status, fault, the session's end: + # ended, revoked, left_for_drain or left_to_master) + "written": (None, None, "ready", None, "ended"), + "short_write": (5, None, "fault", "put ranges: moved 5 of 8 bytes", "revoked"), + # A TCP target may still write into the object's space. + "transfer_fail": (-800, None, "fault", "put ranges: TRANSFER_FAIL(-800)", "left_to_master"), + "session_lost": (-600, None, "fault", "put ranges: INVALID_PARAMS(-600)", "revoked"), + "end_refused": (None, -704, "fault", "put session end: OBJECT_NOT_FOUND(-704)", "revoked"), + "end_not_owner": (None, -601, "fault", "put session end: ILLEGAL_CLIENT(-601)", "revoked"), + # The end may still reach the master. + "end_timed_out": (None, -901, "fault", "put session end: RPC_TIMEOUT(-901)", "left_for_drain"), +} + + +@pytest.mark.parametrize("case", PUT_WRITE_ANSWERS) +def test_put_ranges_and_end_answers_map_to_row_statuses(case: str) -> None: + ranges, end, status, fault, session = PUT_WRITE_ANSWERS[case] + bench = Bench() + bench.write(0, b"aaaaaaaa") + bench.fake.script(PUT_RANGES, [ranges]) + bench.fake.script(PUT_END, [end]) + bench.fake.script(PUT_END, [end]) # the retry of an unanswered end + results, access_ended = bench.run("put", "A") + assert summary(results) == [(status, fault)] + assert access_ended + assert (bench.fake.pool.objects.get(key("A")) == b"aaaaaaaa") == (status == "ready") + assert bench.keys_of(PUT_REVOKE) == ([key("A")] if session == "revoked" else []) + stats = bench.driver.stats() + assert stats["put_sessions_left_for_drain"] == (session == "left_for_drain") + assert stats["put_sessions_left_to_master"] == (session == "left_to_master") + + +EXISTENCE_ANSWERS = { + # batch_is_exist's step: the results of the two rows that already existed + "committed_and_not": ([1, 0], [("ready", None), ("not_ready", None)]), + "check_failed": ([1, -900], [("ready", None), ("fault", "existence check: RPC_FAIL(-900)")]), + # A metadata call touches no row memory: its raise fails only its rows. + "check_raised": (RuntimeError("master gone"), [("fault", METADATA_CALL_FAILED)] * 2), +} + + +@pytest.mark.parametrize("case", EXISTENCE_ANSWERS) +def test_only_rows_that_already_exist_are_checked_for_existence(case: str) -> None: + exist, expected = EXISTENCE_ANSWERS[case] + bench = Bench() + # A starts, B and C exist (OBJECT_ALREADY_EXISTS), D finds the pool full. + bench.fake.script(PUT_START, [None, -705, -705, -200]) + bench.fake.script("batch_is_exist", exist) + results, access_ended = bench.run("put", "A", "B", "C", "D") + assert summary(results) == [("ready", None), *expected, ("declined", None)] + assert bench.keys_of("batch_is_exist") == [key("B"), key("C")] + assert access_ended + + +# Session bookkeeping. + + +def reply_lost(answers): + """The end committed every object and dropped the sessions, but the reply never came.""" + raise RuntimeError("reply lost") + + +class SessionScenario(NamedTuple): + """A put of rows A, B and C with some SDK calls scripted, and what must come of its sessions.""" + + steps: dict # method -> one FakeMooncakeStore.script step per call, in call order + rows: str # each row ready (R) or faulted (F) + access_ends: bool = True + revoked: str = "" # rows whose session the put revokes, or tries to + left_for_drain: str = "" # rows whose session is left open for drain to revoke + left_to_master: str = "" # rows whose session is left open for the master alone to free + + +SESSION_SCENARIOS = { + "all_written": SessionScenario({}, "RRR"), + # A TCP target may still write into B's object space: only the master may free it. + "transfer_fail": SessionScenario({PUT_RANGES: [[None, -800, None]]}, "RFR", left_to_master="B"), + "short_write": SessionScenario({PUT_RANGES: [[None, 5, None]]}, "RFR", revoked="B"), + "end_retried": SessionScenario({PUT_END: [[-900, None, None], [None]]}, "RRR"), + # Ends that may still reach the master. + "end_fails_twice": SessionScenario( + {PUT_END: [[-901, None, None], [-900]]}, "FRR", left_for_drain="A" + ), + "end_retry_refused": SessionScenario( + {PUT_END: [[-901, None, None], [-704]]}, "FRR", left_for_drain="A" + ), + "end_retry_raises": SessionScenario( + {PUT_END: [[-901, None, None], RuntimeError("x")]}, "FRR", left_for_drain="A" + ), + "end_raises": SessionScenario( + {PUT_END: [RuntimeError("rpc lost")]}, "FFF", left_for_drain="ABC" + ), + "end_committed_then_raised": SessionScenario( + {PUT_END: [reply_lost]}, "FFF", left_for_drain="ABC" + ), + "end_and_revoke_fail": SessionScenario( + {PUT_END: [[-704, None, None]], PUT_REVOKE: [[-900]]}, + "FRR", + revoked="A", + left_for_drain="A", + ), + # The SDK may have opened any of the sessions. + "start_raises": SessionScenario( + {PUT_START: [RuntimeError("rpc lost")]}, "FFF", left_for_drain="ABC" + ), + "start_miscounts": SessionScenario( + {PUT_START: [lambda answers: answers[:-1]]}, "FFF", left_for_drain="ABC" + ), + # The scatter may still be running: nothing is ended or revoked. + "ranges_raise": SessionScenario( + {PUT_RANGES: [RuntimeError("bad pointer")]}, "FFF", access_ends=False, left_to_master="ABC" + ), + "ranges_miscount": SessionScenario( + {PUT_RANGES: [lambda answers: answers[:-1]]}, "FFF", access_ends=False, left_to_master="ABC" + ), +} + + +@pytest.mark.parametrize("name", SESSION_SCENARIOS) +def test_every_started_put_session_is_ended_revoked_or_left_open(name: str) -> None: + scenario = SESSION_SCENARIOS[name] + bench = Bench() + for slot, data in enumerate((b"aaaaaaaa", b"bbbbbbbb", b"cccccccc")): + bench.write(slot, data) + for method, method_steps in scenario.steps.items(): + for step in method_steps: + bench.fake.script(method, step) + results, access_ended = bench.run("put", "A", "B", "C") + assert "".join(r.status.value[0].upper() for r in results) == scenario.rows + assert access_ended is scenario.access_ends + + assert bench.keys_of(PUT_REVOKE) == [key(n) for n in scenario.revoked] + settled = Counter(bench.fake.keys_answered(PUT_END, 0) + bench.keys_of(PUT_REVOKE)) + assert set(settled.values()) <= {1}, f"a session ended or revoked twice: {settled}" + for_drain = {key(n) for n in scenario.left_for_drain} + to_master = {key(n) for n in scenario.left_to_master} + stats = bench.driver.stats() + assert stats["put_sessions_left_for_drain"] == len(for_drain) + assert stats["put_sessions_left_to_master"] == len(to_master) + # Every session the SDK still holds is left open; the SDK may have dropped one the driver + # could not see end, as when an end landed but its reply did not. + assert bench.fake.put_sessions <= for_drain | to_master + # Only a row written whole may be ended: its object becomes readable as it stands. + written = bench.fake.keys_answered(PUT_RANGES, SLOT) + assert set(bench.keys_of(PUT_END)) <= set(written) + ready = [key(n) for n, r in zip("ABC", results) if r.status is RowStatus.READY] + assert set(ready) == set(bench.fake.keys_answered(PUT_END, 0)) + # A drain revokes the sessions left for it, never those left to the master. + if scenario.access_ends: + assert bench.driver.drain(WAIT_S) is True + assert bench.driver.stats()["put_sessions_left_for_drain"] == 0 + assert bench.driver.stats()["put_sessions_left_to_master"] == len(to_master) + assert bench.fake.put_sessions == to_master + assert not to_master & set(bench.keys_of(PUT_REVOKE)) + + +def test_sessions_left_open_are_revoked_only_in_drain() -> None: + bench = Bench() + bench.write(0, b"aaaaaaaa") + bench.fake.script(PUT_END, [-704]) + bench.fake.script(PUT_REVOKE, [-900]) + bench.run("put", "A") + assert bench.driver.stats()["put_sessions_left_for_drain"] == 1 + bench.write(1, b"bbbbbbbb") + bench.store(C=b"cccccccc") + bench.run("put", "X", "B") + bench.run("get", "C") + assert bench.keys_of(PUT_REVOKE) == [key("A")] # only its own first try + results, _ = bench.run("put", "A") # the driver holds A's session: it starts no other + assert summary(results) == [("not_ready", None)] # and no complete object exists + assert bench.keys_of(PUT_START).count(key("A")) == 1 + + bench.fake.script(PUT_REVOKE, [-900]) + assert bench.driver.drain(WAIT_S) is True # sessions left open hold no lender memory + assert bench.keys_of(PUT_REVOKE) == [key("A")] * 2 + assert bench.driver.stats()["put_sessions_left_for_drain"] == 1 # retried by the next drain + assert bench.driver.drain(0.0) is True + assert bench.driver.stats()["put_sessions_left_for_drain"] == 0 + assert bench.fake.put_sessions == set() + + +def test_a_put_whose_worker_thread_exits_leaves_its_session_open_for_drain( + thread_escapes: list[BaseException], +) -> None: + bench = Bench() + bench.write(0, b"aaaaaaaa") + bench.fake.script(PUT_END, SystemExit()) # escapes every handler: the worker thread ends + results, access_ended = bench.run("put", "A") + assert summary(results) == [("fault", "the worker stopped before the row was decided")] + assert not access_ended + assert key("A") in bench.fake.put_sessions + assert bench.driver.stats()["put_sessions_left_for_drain"] == 1 + results, _ = bench.run("put", "A") # the driver holds A's session: it starts no other + assert summary(results) == [("not_ready", None)] + assert bench.keys_of(PUT_START).count(key("A")) == 1 + wait_for_an_escape(thread_escapes) + + +def test_a_key_another_put_of_this_driver_holds_is_not_ready_without_asking_the_sdk() -> None: + bench = Bench() + bench.fill_sources("put", "A", "B") + gate = bench.fake.block(PUT_RANGES) + first = bench.driver.put(bench.rows("A")) + try: + assert gate.entered.wait(WAIT_S) # A's session is open, its scatter held + results, _ = bench.settle(bench.driver.put(bench.rows("A", "B"))) + assert summary(results) == [("not_ready", None), ("ready", None)] + results, _ = bench.settle(bench.driver.put(bench.rows("A"))) # nothing left to start + assert summary(results) == [("not_ready", None)] + assert [call.args[0] for call in bench.fake.calls_to(PUT_START)] == [[key("A")], [key("B")]] + finally: + gate.release.set() + bench.settle(first) + assert bench.fake.pool.objects[key("A")] == b"A" * SLOT + assert skipped_counts(bench) == {"rows_not_ready_held_by_own_put": 2} + results, _ = bench.run("put", "A") # A's put ended, so the key is free again + assert bench.keys_of(PUT_START).count(key("A")) == 2 + assert summary(results) == [("ready", None)] # it exists already, committed + + +@pytest.mark.parametrize("left", ["for_drain", "to_master"]) +def test_a_key_whose_session_was_left_open_is_ready_once_its_object_exists(left: str) -> None: + bench = Bench() + bench.fill_sources("put", "A") + if left == "to_master": + bench.fake.script(PUT_RANGES, [-800]) + else: # the end and its retry go unanswered + bench.fake.script(PUT_END, [-901]) + bench.fake.script(PUT_END, [-901]) + bench.run("put", "A") + assert bench.driver.stats()[f"put_sessions_left_{left}"] == 1 + results, _ = bench.run("put", "A") + assert summary(results) == [("not_ready", None)] # no complete object yet + # The master freed the unfinished put, and another client then wrote the key whole. + bench.fake.pool.reserved.pop(key("A")) + bench.store(A=b"a" * SLOT) + results, _ = bench.run("put", "A") + assert summary(results) == [("ready", None)] # whoever wrote it + assert bench.keys_of(PUT_START) == [key("A")] # its session was never started again + assert bench.keys_of("batch_is_exist") == [key("A")] * 2 + + +def test_close_needs_a_drain_that_returned_true() -> None: + bench = Bench() + bench.fill_sources("put", "A") + gate = bench.fake.block(PUT_RANGES) + put = bench.driver.put(bench.rows("A")) + try: + assert gate.entered.wait(WAIT_S) + assert bench.driver.drain(0.0) is False # the put is still running + finally: + gate.release.set() + bench.settle(put) + bench.driver.unregister(bench.handle) + with pytest.raises(DriverRejected, match="drain that returned True"): + bench.driver.close() + assert bench.fake.calls_to("close") == [] + assert bench.driver.drain(WAIT_S) is True + bench.driver.close() + assert len(bench.fake.calls_to("close")) == 1 + + +def test_drain_leaves_the_session_of_an_unprovable_put_to_the_master() -> None: + bench = Bench() + bench.fake.script(PUT_RANGES, RuntimeError("bad pointer")) + _, access_ended = bench.run("put", "A") + assert not access_ended + assert bench.driver.drain(0.0) is False + assert bench.fake.calls_to(PUT_REVOKE) == [] + assert bench.driver.stats()["put_sessions_left_to_master"] == 1 + + +LEASE_EXPIRED_CASES = { + # the first read's answers, the reread's, the results, whether access ended + "reread_ready": ([None, -707, None], [None], "RRR", True), + "expired_again": ([None, -707, None], [-707], "RFR", True), + "first_read_failed": ([None, -800, -707], [None], "RFR", False), + "reread_failed": ([None, None, -707], [-800], "RRF", False), + # The first read already wrote the slot, so the reread's answer cannot make it a MISS or + # NOT_READY. + "reread_not_found": ([None, -707, None], [-704], "RFR", True), + "reread_not_ready": ([None, -707, None], [-703], "RFR", True), +} + + +@pytest.mark.parametrize("case", LEASE_EXPIRED_CASES) +def test_rows_whose_lease_expired_are_reread_once_and_both_reads_must_end_access(case: str) -> None: + first, reread, statuses, ended = LEASE_EXPIRED_CASES[case] + bench = Bench() + bench.fill_sources("get", "A", "B", "C") + bench.fake.script(GET, first) + bench.fake.script(GET, reread) + results, access_ended = bench.run("get", "A", "B", "C") + assert "".join(r.status.value[0].upper() for r in results) == statuses + expired = [key(name) for name, answer in zip("ABC", first) if answer == -707] + assert [call.args[0] for call in bench.fake.calls_to(GET)] == [[key(n) for n in "ABC"], expired] + assert access_ended is ended + assert bench.driver.drain(WAIT_S) is ended # False at once when unprovable + + +def test_an_object_being_written_is_not_ready_for_put_and_get(monkeypatch) -> None: + warnings: list[str] = [] + monkeypatch.setattr(mooncake_module.logger, "warning", warnings.append) + pool = FakePool() + writer = FakeMooncakeStore(pool) + assert writer.batch_put_session_start([key("K")], [SLOT]) == [0] # another client is writing + bench = Bench(pool=pool) + results, access_ended = bench.run("get", "K") + assert summary(results) == [("not_ready", None)] and access_ended + results, access_ended = bench.run("put", "K") + assert summary(results) == [("not_ready", None)] and access_ended + stats = bench.driver.stats() + assert stats["rows_not_ready_no_complete_replica"] == stats["rows_not_ready_uncommitted"] == 1 + # Each part logs once it ended, a moment after its operation settled. + for reason in ("no_complete_replica", "uncommitted"): + wait_until( + lambda: any(key("K") in line and f"not_ready ({reason})" in line for line in warnings), + f"the {reason} warning", + ) + + +# Access end. + + +FAILURES = { + "raises": RuntimeError("rpc lost"), + "raises_value_error": ValueError("bad pointer"), + "short": lambda answers: answers[:-1], + "long": lambda answers: [*answers, 0], + "empty": lambda answers: [], +} + + +@pytest.mark.parametrize("failure", FAILURES) +@pytest.mark.parametrize("call", ["put", "get", "get_reread"]) +def test_access_stays_open_when_a_data_call_raises_or_miscounts(call: str, failure: str) -> None: + bench = Bench() + direction = "put" if call == "put" else "get" + bench.fill_sources(direction, "A", "B") + if call == "get_reread": + bench.fake.script(GET, [None, -707]) + bench.fake.script(DATA_CALL[direction], FAILURES[failure]) + results, access_ended = bench.run(direction, "A", "B") # nothing escapes the driver + assert not access_ended and bench.driver.access_unprovable(bench.last) + unproven = results[1:] if call == "get_reread" else results + assert all(r.status is RowStatus.FAULT and " raised " in r.fault for r in unproven) + if call == "get_reread": + assert results[0].status is RowStatus.READY # read whole the first time + assert bench.driver.stats()["ops_unprovable"] == 1 + assert bench.driver.drain(0.0) is False + + +@pytest.mark.parametrize("direction", ["put", "get"]) +def test_a_fault_past_the_sdk_calls_keeps_access_open(direction: str, monkeypatch) -> None: + def broken(*_): + raise ValueError("mapping bug") + + bench = Bench() + if direction == "get": + monkeypatch.setattr(mooncake_module, "_result_of_get", broken) + bench.fake.script(GET, [5]) # a row read short goes past the all-whole shortcut + else: + monkeypatch.setattr(MooncakeDriver, "_end_sessions", broken) + bench.fill_sources(direction, "A") + results, access_ended = bench.run(direction, "A") + assert summary(results) == [("fault", f"{direction} failed: ValueError('mapping bug')")] + assert not access_ended + # A put's session may still be open: the driver keeps holding the key. + assert bench.driver.stats()["put_sessions_left_for_drain"] == (direction == "put") + + +# Admission, registration, concurrency. + + +def test_bad_rows_are_rejected_before_the_sdk() -> None: + bench = Bench() + with pytest.raises(DriverRejected, match="at most"): + bench.driver.put(RowBatch.of([RowIO(key("A"), (Span(bench.buffer.address, 1 << 31),))])) + two_spans = (bench.buffer.span(0, SLOT), bench.buffer.span(2 * SLOT, SLOT)) + with pytest.raises(DriverRejected, match="one span"): + bench.driver.put(RowBatch.of([RowIO(key("A"), two_spans)])) + assert [c for c in bench.fake.calls if c.method.startswith("batch_")] == [] + + +def test_a_batch_whose_extent_leaves_its_registration_is_rejected_before_the_sdk() -> None: + bench = Bench() + other = HostBuffer(SLOT) # a second registration, apart from the first + bench.driver.register(other.address, SLOT, RegionLayout(SLOT, 1)) + first, _ = bench.rows("A", "B") + with pytest.raises(DriverRejected, match="outside every live registration"): + bench.driver.get(RowBatch.of([first, RowIO(key("C"), (other.span(),))])) # gap: no row + assert [c for c in bench.fake.calls if c.method.startswith("batch_")] == [] + results, access_ended = bench.run("get", "A", "B") # one registration holds them + assert summary(results) == [("miss", None), ("miss", None)] and access_ended + + +def test_an_empty_batch_is_rejected_before_the_sdk() -> None: + bench = Bench() + empty = RowBatch.of_slots( + "k/", np.zeros((0, 2), np.uint8), np.zeros(0, np.int64), SLOT, bench.buffer.span() + ) + with pytest.raises(DriverRejected, match="at least one row"): + bench.driver.get(empty) + assert bench.driver.drain(WAIT_S) is True # nothing was left in flight + + +def test_a_part_whose_sdk_lists_cannot_be_built_faults_and_ends_its_access( + monkeypatch: pytest.MonkeyPatch, +) -> None: + bench = Bench() + + def broken(*_) -> list[str]: + raise ValueError("a key bug") + + monkeypatch.setattr(RowBatch, "keys", broken) + results, access_ended = bench.run("get", "A") + assert summary(results) == [ + ("fault", "get failed before any SDK call: ValueError('a key bug')") + ] + assert access_ended and not bench.driver.access_unprovable(bench.last) + assert bench.fake.calls_to(GET) == [] + assert bench.driver.drain(WAIT_S) is True # no operation turned unprovable + + +def test_register_and_unregister_follow_the_protocol() -> None: + fake = FakeMooncakeStore() + driver = MooncakeDriver(fake) + region, other = HostBuffer(64), HostBuffer(64) + handle = driver.register(region.address, 64, RegionLayout(SLOT, 8)) + with pytest.raises(DriverRejected, match="layout needs"): + driver.register(other.address, SLOT, RegionLayout(SLOT, 2)) + assert len(fake.calls_to("register_buffer")) == 1 + + fake.script("register_buffer", [-1]) + with pytest.raises(RegistrationError, match="INTERNAL_ERROR") as refused: + driver.register(other.address, 64, RegionLayout(SLOT, 8)) + assert refused.value.handle is None + fake.script("register_buffer", RuntimeError("rpc lost")) + with pytest.raises(RegistrationError, match="rpc lost") as unknown: + driver.register(other.address, 64, RegionLayout(SLOT, 8)) + assert unknown.value.handle is not None + with pytest.raises(StoreError): # the SDK never kept it, so it cannot revoke it either + driver.unregister(unknown.value.handle) + + def revocations() -> int: + return [call.args for call in fake.calls_to("unregister_buffer")].count((region.address,)) + + fake.script("unregister_buffer", [-1]) + with pytest.raises(StoreError, match="INTERNAL_ERROR"): + driver.unregister(handle) + with pytest.raises(DriverRejected, match="overlaps"): # the failed revocation left it live + driver.register(region.address + SLOT, SLOT, RegionLayout(SLOT, 1)) + driver.unregister(handle) # and retryable + assert revocations() == 2 + driver.unregister(handle) # and is idempotent once it succeeded + assert revocations() == 2 + + +def wait_until_draining(driver: MooncakeDriver) -> None: + """Returns once a drain waits: admission, under the drain's lock, refuses a stray row as + draining rather than as outside every region. Nothing reaches the SDK.""" + stray = HostBuffer(SLOT) + deadline = time.monotonic() + WAIT_S + while time.monotonic() < deadline: + with pytest.raises(DriverRejected) as refused: + driver.get(RowBatch.of([RowIO(key("probe"), (stray.span(),))])) + if "draining" in str(refused.value): + return + pytest.fail("no drain started") + + +def test_unregister_is_refused_and_drain_waits_while_a_put_runs() -> None: + bench = Bench() + bench.fill_sources("put", "A") + gate = bench.fake.block(PUT_RANGES) + putter = threading.Thread(target=lambda: bench.run("put", "A")) + drained: list[bool] = [] + drainer = threading.Thread(target=lambda: drained.append(bench.driver.drain(WAIT_S))) + putter.start() + try: + assert gate.entered.wait(WAIT_S) # inside the scatter + with pytest.raises(DriverRejected, match="may still access"): + bench.driver.unregister(bench.handle) + drainer.start() + wait_until_draining(bench.driver) + assert drainer.is_alive() # still running, not unprovable: the drain waits + finally: + released = time.monotonic() + gate.release.set() + putter.join(WAIT_S) + drainer.join(WAIT_S) + assert drained == [True] + assert time.monotonic() - released < PROMPT_S # the put's end woke the drain + bench.driver.unregister(bench.handle) + assert len(bench.fake.calls_to("unregister_buffer")) == 1 + + +def test_an_unprovable_operation_pins_its_region_and_fails_drain_at_once() -> None: + bench = Bench() + bench.fill_sources("get", "A", "B") + bench.fake.script(GET, [-800]) + bench.run("get", "A") + gate = bench.fake.block(GET) + blocked = bench.driver.get(bench.rows("A", "B")[1:]) + try: + assert gate.entered.wait(WAIT_S) # another get is still running + started = time.monotonic() + assert bench.driver.drain(WAIT_S) is False # waits for no running call + assert time.monotonic() - started < PROMPT_S + finally: + gate.release.set() + bench.settle(blocked) + with pytest.raises(DriverRejected, match="may still access"): + bench.driver.unregister(bench.handle) + with pytest.raises(DriverRejected, match="drain"): + bench.driver.close() + assert bench.fake.calls_to("close") == [] + + +def test_close_needs_a_drain_closes_the_sdk_client_once_and_then_queries_raise() -> None: + bench = Bench() + bench.fill_sources("put", "A") + bench.run("put", "A") + with pytest.raises(DriverRejected, match="drain"): + bench.driver.close() + assert bench.driver.drain(WAIT_S) is True + bench.driver.unregister(bench.handle) + bench.driver.close() + bench.driver.close() # idempotent + assert len(bench.fake.calls_to("close")) == 1 + with pytest.raises(RuntimeError, match="closed"): + bench.driver.poll(bench.last) + with pytest.raises(RuntimeError, match="closed"): + bench.driver.drain(0.0) + + +def test_hints_and_foreign_handles_are_programming_errors() -> None: + bench, other = Bench(), Bench() + with pytest.raises(DriverRejected, match="hints"): + bench.driver.get(bench.rows("A"), "a peer") + other.fill_sources("put", "A") + other.run("put", "A") + for query in ("poll", "access_ended", "access_unprovable", "cancel"): + with pytest.raises(ValueError, match="does not belong"): + getattr(bench.driver, query)(other.last) + with pytest.raises(ValueError, match="does not belong"): + bench.driver.unregister(other.handle) + assert [c for c in bench.fake.calls if c.method.startswith("batch_")] == [] + + +def test_shutdown_ends_the_worker_threads_even_with_an_unprovable_get() -> None: + pool = FakePool() + fake = FakeMooncakeStore(pool) + reader = Endpoint(MooncakeDriver(fake), ["A"]) + pool.objects[store_key("A")] = b"aaaaaaaa" + fake.script(GET, [-800]) # TRANSFER_FAIL: the read may still run, so its access is unproven + attempt = reader.backend.fetch(reader.view("A"), None) + reader.wait_outcome(attempt) + assert reader.backend.shutdown(WAIT_S) == frozenset({attempt.label, "region:0"}) + driver = reader.driver.inner + wait_until(lambda: driver._workers_alive == 0, "the driver's worker threads to exit") + + +def test_backend_round_trip_over_the_fake_store() -> None: + pool = FakePool() # two clients of one master + writer = Endpoint(MooncakeDriver(FakeMooncakeStore(pool)), ["A", "B"]) + reader = Endpoint(MooncakeDriver(FakeMooncakeStore(pool)), ["A", "B", "C", "D"], fill=0x2D) + assert isinstance(writer.backend, Backend) + assert writer.backend.capabilities == Capabilities(fetch=True, publish=True, cancel=True) + writer.write("A", b"aaaaaaaa") + writer.write("B", b"bbbbbbbb") + pool.objects[store_key("D")] = b"short" + + published = writer.backend.publish(writer.view("A", "B"), None) + assert isinstance(published, Attempt) + assert writer.wait_outcome(published).delivered_masks[0].tolist() == [True, True] + assert pool.objects[store_key("A")] == b"aaaaaaaa" + writer.wait_quiesced(published) + + view = reader.view("A", "B", "C") + fetched = reader.backend.fetch(view, None) + outcome = reader.wait_outcome(fetched) + outcome.check_mask_shapes(view) + assert outcome.fault is None + assert outcome.delivered_masks[0].tolist() == [True, True, False] + assert [reader.read(name) for name in "ABC"] == [b"aaaaaaaa", b"bbbbbbbb", b"-" * SLOT] + reader.wait_quiesced(fetched) + + wrong = reader.backend.fetch(reader.view("D"), None) + outcome = reader.wait_outcome(wrong) + assert outcome.delivered_masks[0].tolist() == [False] + assert "object is 5 bytes, expected 8" in outcome.fault + reader.wait_quiesced(wrong) + assert reader.region.buffer.canaries_intact() + writer.close() + reader.close() + + +@pytest.fixture +def work_per_row_or_thread(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, int]]: + """Every key or row a ``RowBatch`` builds and every thread started, each with the thread that + did it.""" + done: list[tuple[str, int]] = [] + keys_iter, key_at = store_module._HexKeys.__iter__, store_module._HexKeys.__getitem__ + row_at, thread_start = RowBatch.__getitem__, threading.Thread.start + + def record(what: str) -> None: + done.append((what, threading.get_ident())) + + def every_key(self): + record("keys") + return keys_iter(self) + + def one_key(self, index): + if not isinstance(index, slice): + record("key") + return key_at(self, index) + + def one_row(self, index): + if not isinstance(index, slice): + record("row") + return row_at(self, index) + + def start(self): + record("thread") + return thread_start(self) + + monkeypatch.setattr(store_module._HexKeys, "__iter__", every_key) + monkeypatch.setattr(store_module._HexKeys, "__getitem__", one_key) + monkeypatch.setattr(RowBatch, "__getitem__", one_row) + monkeypatch.setattr(threading.Thread, "start", start) + return done + + +def test_fetch_and_publish_build_no_key_or_row_and_start_no_thread_on_the_calling_thread( + work_per_row_or_thread: list[tuple[str, int]], +) -> None: + # The calling thread is the executor loop's: it builds no store key or RowIO and starts no + # thread; keys and rows are built on the driver's threads. Built without the test endpoint, + # whose driver spy reads every row it is handed. + rows = 64 + region = CpuRegion(SLOT, 2 * rows) + binding = RegionBinding((region,), MooncakeDriver(FakeMooncakeStore()), FakeHold()) + binding.register() + wake = WakeCounter() + backend = BlobBackend(binding, wake, BlobConfig(key_prefix="t")) + keys = np.stack([object_key(f"R{i}") for i in range(rows)]) + for i in range(rows): + region.write(i, f"R{i}".encode().ljust(SLOT, b".")) + offered = CpuView([Rows(0, keys, region.slots(*range(rows)), 0)]) + wanted = CpuView([Rows(0, keys.copy(), region.slots(*range(rows, 2 * rows)), 0)]) + caller = threading.get_ident() + for submit, view in ((backend.publish, offered), (backend.fetch, wanted)): + work_per_row_or_thread.clear() + attempt = submit(view, None) + assert [what for what, thread in work_per_row_or_thread if thread == caller] == [] + wake.wait_for(lambda: attempt.poll() is not None and attempt.quiesced(), "the attempt") + assert attempt.poll().delivered_masks[0].all() + assert region.read(rows + 7) == b"R7......" + assert backend.shutdown(WAIT_S) == frozenset() From ec7bd8d3742e773f9f48029146f6ee26f547a828 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Tue, 27 Oct 2026 08:23:08 +0000 Subject: [PATCH 02/10] [None][feat] Open a Mooncake store driver from a store config A deployment describes its store in a config file of its own, while the backend needs a MooncakeDriver over a set-up SDK client and a BlobConfig. mooncake_setup turns the one into the other. load_store_config reads the store config, a YAML file or the mapping parsed from one: a blob section for the caller's reader (BlobConfig for the blob backend), and a mooncake section read into MooncakeConfig, which also sets the driver's workers, rows_per_call and bytes_per_call (defaults measured on GB200 over RDMA). Unknown sections and keys, and values of the wrong type, are refused with their section; so are an unreadable file, a YAML error and a master address that is not one host:port. open_sdk refuses TCP's legacy framing, the local hot cache and the newer transfer engine (MC_USE_TENT, MC_USE_TEV1), whose behaviour the driver's rules were not read for; logs the Mooncake environment before setup; imports mooncake.store only when called, and refuses it unless every installed distribution that lists the module is a pinned 0.3.13 build; waits for the master; resolves the RDMA devices and the local address; and closes a client whose setup failed. open_mooncake_driver reads both sections before it opens the SDK, closes the client again when the driver cannot start its workers, and returns the driver with what the caller's reader returned. A client in service is never closed here. test_mooncake_setup.py covers config errors, setup arguments, device detection, a master that comes up late, the refused settings and SDK builds, a missing SDK, a half-open client, the opened driver working through its client, and the blob package's import layering. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../_torch/disaggregation/blob/README.md | 3 + .../_torch/disaggregation/blob/__init__.py | 7 +- .../disaggregation/blob/mooncake_setup.py | 412 +++++++++++++ .../blob/test_mooncake_setup.py | 572 ++++++++++++++++++ 4 files changed, 993 insertions(+), 1 deletion(-) create mode 100644 tensorrt_llm/_torch/disaggregation/blob/mooncake_setup.py create mode 100644 tests/unittest/_torch/disaggregation/blob/test_mooncake_setup.py diff --git a/tensorrt_llm/_torch/disaggregation/blob/README.md b/tensorrt_llm/_torch/disaggregation/blob/README.md index 612437cc5695..ae06c7aba0d0 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/README.md +++ b/tensorrt_llm/_torch/disaggregation/blob/README.md @@ -33,6 +33,7 @@ store SDK | `store.py` | the driver boundary: `StoreDriver`, the optional `RegistersRegionsTogether`, `RowBatch`, `RowIO`, `OnChange`, `RowStatus`, `DriverCapabilities`, `DriverRejected`, `RegistrationError` | | `keys.py` | the store key encoding | | `mooncake.py` | `MooncakeDriver`, a `StoreDriver` over a set-up Mooncake 0.3.13 SDK client | +| `mooncake_setup.py` | `open_mooncake_driver`: reads a store config and opens the SDK client and the driver | | `../testing/` | what a driver's own CI imports from the wheel: the conformance suite, the in-memory store and drivers, the broken drivers that prove every check bites | ## Running the tests @@ -55,6 +56,8 @@ All of it is CPU-only. The files, top down: - `test_testing_isolation.py`: no engine module imports the test tooling. - `test_mooncake_driver.py`: the Mooncake driver over a fake SDK (`fake_mooncake.py`) that answers as 0.3.13 does, the conformance suite included. +- `test_mooncake_setup.py`: the store config, the SDK's setup and the refused settings and builds, + and that each module imports only its layer. ## Writing a driver diff --git a/tensorrt_llm/_torch/disaggregation/blob/__init__.py b/tensorrt_llm/_torch/disaggregation/blob/__init__.py index 7809fd32d115..56e778597f47 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/blob/__init__.py @@ -13,7 +13,12 @@ public contract, and the store key encoding. How an object is laid out in a store is its driver's private format. - ``mooncake``: ``MooncakeDriver``, a ``StoreDriver`` over a set-up Mooncake SDK client. +- ``mooncake_setup``: reads a store config and opens the SDK client; ``open_mooncake_driver`` + reads the config's ``blob`` section with the caller's reader (``BlobConfig`` for the backend) + before it opens the SDK, and returns the driver with what the reader returned. -No module imports torch, ``tensorrt_llm._torch.pyexecutor`` or a store SDK at import time. +No module imports torch, ``tensorrt_llm._torch.pyexecutor`` or a store SDK at import time; only +``mooncake_setup.open_sdk`` imports ``mooncake.store``, when called. Each module's imports are +checked by ``test_mooncake_setup.py::test_blob_modules_import_only_their_layer``. ``README.md`` next to this file says how to write a driver and how to assemble a backend. """ diff --git a/tensorrt_llm/_torch/disaggregation/blob/mooncake_setup.py b/tensorrt_llm/_torch/disaggregation/blob/mooncake_setup.py new file mode 100644 index 000000000000..b09acd306ecc --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/blob/mooncake_setup.py @@ -0,0 +1,412 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Opening Mooncake from a store config: the deployment glue around the driver. + +A store config has two sections: ``blob``, read by the caller's reader, and ``mooncake``, read +here; both are checked before the SDK is opened. Mooncake reads its own environment variables, some +once per process (``MOONCAKE_STORE_CHECKSUM``), some per client (``MC_STORE_MEMCPY``, +``MC_USE_TENT``) and ``MC_TCP_PROTO`` on every transfer; this module logs them before setup and +never sets them. ``open_sdk`` refuses the settings the driver's rules do not cover: TCP's legacy +framing (``MC_TCP_PROTO=1``), the local hot cache (``MC_STORE_LOCAL_HOT_CACHE_SIZE``), the newer +transfer engine (``MC_USE_TENT``, ``MC_USE_TEV1``) and any SDK build other than the pinned ones. + +Over RDMA, a process that used ibverbs before its first Mooncake client (importing +``tensorrt_llm`` does) logs ``RDMA context setup failed: fork compatibility: Invalid argument`` +once: ``ibv_fork_init`` only works before any other verbs call. It is harmless on Linux 5.13 and +later, which copy pinned pages on fork, and the devices come up normally. +""" + +from __future__ import annotations + +import importlib.metadata +import os +import pathlib +import re +import socket +import sys +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Any, Literal, TypeVar + +import yaml + +from tensorrt_llm.logger import logger + +from .mooncake import DEFAULT_BYTES_PER_CALL, DEFAULT_ROWS_PER_CALL, DEFAULT_WORKERS, MooncakeDriver + +__all__ = [ + "MooncakeConfig", + "load_store_config", + "open_mooncake_driver", + "open_sdk", +] + +RECOMMENDED_SDK_DISTRIBUTION = "mooncake-transfer-engine-cuda13" # the one TensorRT-LLM installs +# The builds of the SDK with the transports the driver's rules were read for (RDMA and TCP). +PINNED_SDK_DISTRIBUTIONS = ( + RECOMMENDED_SDK_DISTRIBUTION, + "mooncake-transfer-engine", + "mooncake-transfer-engine-non-cuda", +) +PINNED_SDK_VERSION = "0.3.13" +IB_SYSFS_ROOT = "/sys/class/infiniband" +MASTER_WAIT_S = 60.0 +MASTER_RETRY_S = 0.5 +SECTIONS = ("blob", "mooncake") +# Logged before setup, as Mooncake will read them. +_LOGGED_ENVIRONMENT = ( + "MC_STORE_MEMCPY", + "MOONCAKE_STORE_CHECKSUM", + "WITH_NVIDIA_PEERMEM", + "MC_TCP_SLICE_SIZE", + "MC_TCP_PROTO", + "MC_STORE_LOCAL_HOT_CACHE_SIZE", + "MC_USE_TENT", + "MC_USE_TEV1", +) +BlobSectionT = TypeVar("BlobSectionT") +# Reserved for documentation (RFC 5737): a datagram socket connected to it reveals the route off +# this host without sending anything. +_ROUTE_PROBE_ENDPOINT = ("192.0.2.1", 9) + + +@dataclass(frozen=True) +class MooncakeConfig: + """The ``mooncake`` section of a store config. + + Attributes: + master_server_address: The master's ``host:port``. + global_segment_size: Bytes of this process's memory that hold pool objects (its + segment). Required: the objects in it go when the process does; 0 leaves holding + objects to other processes. + protocol: ``"rdma"`` or ``"tcp"``. + device_name: Comma-separated RDMA devices; empty picks every active InfiniBand device at + the highest rate. + local_hostname: The address peers reach this process at; empty takes the address the + route to the master leaves from. + metadata_server: Where segments are described; ``P2PHANDSHAKE`` needs no server. + local_buffer_size: Bytes for the SDK's copying calls, which the driver never makes; 0 is + allowed, and the default is the SDK's. + workers: The most threads the driver runs SDK calls on. One put call does not fill the + links, and each running get busy-polls a core; with two workers or more, puts leave + one to gets. + rows_per_call: The most rows of one SDK call; the driver cuts larger operations into + several. The SDK converts a call's Python lists while holding the GIL, so this bounds + how long one call holds it. + bytes_per_call: The most payload bytes of one SDK call that moves rows. Larger calls return + into Python fewer times per byte, and next to a busy executor thread each return waits + for the GIL; but a call must end well inside the read lease (10 s), and an operation + holds all its slots until its last call returned, or for good once a get answers + TRANSFER_FAIL. + + The defaults of the last three were measured on GB200 over RDMA. + """ + + master_server_address: str + global_segment_size: int + protocol: Literal["rdma", "tcp"] = "rdma" + device_name: str = "" + local_hostname: str = "" + metadata_server: str = "P2PHANDSHAKE" + local_buffer_size: int = 16 << 20 + workers: int = DEFAULT_WORKERS + rows_per_call: int = DEFAULT_ROWS_PER_CALL + bytes_per_call: int = DEFAULT_BYTES_PER_CALL + + def __post_init__(self) -> None: + for name in ("master_server_address", "device_name", "local_hostname", "metadata_server"): + if not isinstance(getattr(self, name), str): + raise ValueError(f"{name} must be a str, got {getattr(self, name)!r}") + if "://" in self.master_server_address: # the waits below need one host:port + raise ValueError( + "master_server_address must be one master's host:port; an HA master list " + f"is not supported, got {self.master_server_address!r}" + ) + if self.protocol not in ("rdma", "tcp"): + raise ValueError(f"protocol must be rdma or tcp, got {self.protocol!r}") + for name, least in ( + ("global_segment_size", 0), + ("local_buffer_size", 0), + ("workers", 1), + ("rows_per_call", 1), + ("bytes_per_call", 1), + ): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, int) or value < least: + raise ValueError(f"{name} must be an int of at least {least}, got {value!r}") + + +def load_store_config( + config: str | os.PathLike | Mapping, +) -> tuple[Mapping[str, object], MooncakeConfig]: + """Reads a store config: a YAML file, or the mapping parsed from one. + + Returns: + The ``blob`` section (empty when absent), and the ``mooncake`` section. + + Raises: + ValueError: The file cannot be read or is not YAML, an unknown section or key, a missing + ``mooncake`` section or required key, or a value of the wrong type. + """ + document = config + if not isinstance(document, Mapping): + try: + with open(config) as file: + document = yaml.safe_load(file) + except (OSError, yaml.YAMLError) as error: + raise ValueError(f"store config {config}: {error}") from error + if not isinstance(document, Mapping): + raise ValueError(f"store config {config} is not a mapping") + unknown = sorted(set(document) - set(SECTIONS), key=str) + if unknown: + raise ValueError(f"store config: unknown sections {unknown}; it has {list(SECTIONS)}") + blob = document.get("blob") + blob = {} if blob is None else blob # ``blob:`` with nothing under it is an empty section + section = document.get("mooncake") + if not isinstance(blob, Mapping) or not isinstance(section, Mapping): + raise ValueError("store config: 'mooncake' is a required mapping, 'blob' an optional one") + try: + return blob, MooncakeConfig(**section) + except (TypeError, ValueError) as error: # an unknown or missing key, or a bad value + raise ValueError(f"store config section 'mooncake': {error}") from error + + +def open_sdk(config: MooncakeConfig) -> Any: + """Opens a set-up ``MooncakeDistributedStore``, on the builder thread: ``setup`` holds the GIL + while it mounts this process's segment. + + Raises: + ValueError: TCP would use its legacy framing, the local hot cache is on, or the newer + transfer engine is selected. + ImportError: ``mooncake.store`` does not import, or is not SDK 0.3.13. + TimeoutError: The master does not accept connections. + RuntimeError: ``setup`` failed; the half-open client is closed. + """ + _refuse_legacy_tcp_framing(config.protocol) + _refuse_hot_cache() + _refuse_tent() + _log_environment() + try: + import mooncake.store as sdk + except ImportError as error: + raise ImportError(_sdk_import_error_message(error)) from error + _require_pinned_builds(sdk) + _wait_for_master(config.master_server_address) + store = sdk.MooncakeDistributedStore() + code = store.setup( + local_hostname=config.local_hostname or _local_address(config.master_server_address), + metadata_server=config.metadata_server, + global_segment_size=config.global_segment_size, + local_buffer_size=config.local_buffer_size, + protocol=config.protocol, + rdma_devices=_rdma_devices(config), + master_server_addr=config.master_server_address, + ) + if code != 0: + store.close() # no registration and no call in flight yet, unlike a client in service + raise RuntimeError( + f"Mooncake setup failed with {code} (master {config.master_server_address}, " + f"protocol {config.protocol})" + ) + return store + + +def open_mooncake_driver( + config: str | os.PathLike | Mapping, + read_blob_section: Callable[[Mapping[str, object]], BlobSectionT], +) -> tuple[MooncakeDriver, BlobSectionT]: + """Opens the Mooncake driver a store config describes, once both its sections are read: opening + waits for the master and mounts this process's segment, so a bad config fails before that. + + Args: + config: A store config, as ``load_store_config`` takes it. + read_blob_section: Reads the ``blob`` section (for the blob backend, into a + ``BlobConfig``); what it returns comes back with the driver. + + Raises: + ValueError: A section is wrong, as ``load_store_config`` or ``read_blob_section`` found. + As ``open_sdk`` otherwise. + """ + blob_section, mooncake_config = load_store_config(config) + try: + blob = read_blob_section(blob_section) + except (TypeError, ValueError) as error: # an unknown key, or a bad value + raise ValueError(f"store config section 'blob': {error}") from error + store = open_sdk(mooncake_config) + try: + driver = MooncakeDriver( + store, + workers=mooncake_config.workers, + rows_per_call=mooncake_config.rows_per_call, + bytes_per_call=mooncake_config.bytes_per_call, + ) + except BaseException: + store.close() # no region was registered and no call made: the client is not in service + raise + return driver, blob + + +def _refuse_legacy_tcp_framing(protocol: str) -> None: + """Under ``MC_TCP_PROTO=1`` a TCP write completes once its last chunk is in the local socket + buffer, so a put could be READY before its bytes landed. The SDK reads the variable on every + transfer: setting it after this check goes unnoticed. The same legacy framing is used toward a + segment whose process runs an SDK older than 0.3.12, which this process cannot check, so every + process of the pool must run 0.3.12 or later (``tcp_transport.cpp:720``, + ``transfer_metadata.cpp:644-646``).""" + if protocol == "tcp" and os.environ.get("MC_TCP_PROTO") == "1": + raise ValueError("MC_TCP_PROTO=1 selects TCP framing that reports writes before they land") + + +def _refuse_hot_cache() -> None: + """The SDK copies the slot a get read into its hot cache before it checks the read lease, and + never drops the copy when the master evicts the object. So a read answered LEASE_EXPIRED, torn + by a reuse of the object's memory, stays cached; once the key is put again, every get of it + here is served those bytes as a whole object. Every value but empty or "0" is refused: the SDK + reads it with ``stoull``, which also turns the cache on for values like "64MB" or " 64".""" + if os.environ.get("MC_STORE_LOCAL_HOT_CACHE_SIZE", "0").strip() not in ("", "0"): + raise ValueError("MC_STORE_LOCAL_HOT_CACHE_SIZE turns on a cache that can serve torn reads") + + +def _refuse_tent() -> None: + """Either variable, set to any value, makes the SDK build its transfer engine on TENT, the newer + engine, whose completion rules the driver was not read against (``mooncake`` module + docstring). The SDK reads them when a client is set up.""" + selected = [name for name in ("MC_USE_TENT", "MC_USE_TEV1") if name in os.environ] + if selected: + raise ValueError( + f"{', '.join(selected)} selects a transfer engine the driver does not cover" + ) + + +def _log_environment() -> None: + values = ", ".join(f"{name}={os.environ.get(name, '')}" for name in _LOGGED_ENVIRONMENT) + logger.info(f"mooncake: environment at setup: {values}") + + +def _sdk_import_error_message(error: ImportError) -> str: + install = f"{RECOMMENDED_SDK_DISTRIBUTION}=={PINNED_SDK_VERSION}" + hint = f"Install {install} and put it before NIXL on the path." + shadow = getattr(sys.modules.get("mooncake"), "__file__", None) + if shadow: + hint = f"The 'mooncake' found is {shadow}; NIXL ships one that shadows the SDK's. {hint}" + return f"cannot import mooncake.store ({error}). {hint}" + + +def _require_pinned_builds(sdk: Any) -> None: + """The driver's code mapping and access-end rules were read from 0.3.13 with the RDMA and TCP + transports; on another version or build a put could return early, and a revoke then free + space a transfer still writes. So the imported ``mooncake.store`` must be shipped by one of + ``PINNED_SDK_DISTRIBUTIONS`` at 0.3.13, not by another build, a source tree or a shadowing + package. Builds installed side by side share the module's path and the file is the last one's, + so every build that lists it must be pinned.""" + builds = _builds_listing_file(sdk) + pinned = {(name, PINNED_SDK_VERSION) for name in PINNED_SDK_DISTRIBUTIONS} + if not builds or not set(builds) <= pinned: + found = ", ".join(f"{name}=={version}" for name, version in builds) + raise ImportError( + f"mooncake.store ({getattr(sdk, '__file__', 'no file')}) comes from " + f"{found or 'no installed distribution'}; the driver follows {PINNED_SDK_VERSION} of " + f"{', '.join(PINNED_SDK_DISTRIBUTIONS)} only. Install " + f"{RECOMMENDED_SDK_DISTRIBUTION}=={PINNED_SDK_VERSION}." + ) + + +def _builds_listing_file(module: Any) -> list[tuple[str, str]]: + """Every installed distribution that lists ``module``'s file, by normalized name and version, + including two of one name, as a reinstall over another version can leave.""" + file = getattr(module, "__file__", None) + if file is None: + return [] + path = pathlib.Path(file).resolve() + builds = [] + for build in importlib.metadata.distributions(): + files = [f for f in build.files or () if f.parts[:1] == ("mooncake",)] + if any(pathlib.Path(build.locate_file(f)).resolve() == path for f in files): + name = build.metadata.get("Name") or "a distribution without a name" + builds.append((re.sub(r"[-_.]+", "-", name).lower(), build.version)) + return builds + + +# Three deployment helpers, with the same rules as the Mooncake connector's +# (pyexecutor/connectors/mooncake_store: master.py and config.py), which this package cannot import. + + +def _rdma_devices(config: MooncakeConfig) -> str: + """``device_name``, or for RDMA left empty every active InfiniBand device at the highest rate: + the compute fabric, without the Ethernet ports that share an HCA.""" + if config.device_name or config.protocol != "rdma": + return config.device_name + detected = _fastest_active_ib_devices(IB_SYSFS_ROOT) + if not detected: + logger.warning(f"mooncake: no active InfiniBand device under {IB_SYSFS_ROOT}") + return ",".join(detected) + + +def _fastest_active_ib_devices(sysfs_root: str) -> list[str]: + try: + devices = sorted(os.listdir(sysfs_root)) + except OSError: + return [] + rates = {} + for device in devices: + port = os.path.join(sysfs_root, device, "ports", "1") + link, state = _read_attribute(port, "link_layer"), _read_attribute(port, "state") + rate = _read_attribute(port, "rate").split() # "400 Gb/sec (4X NDR)" + if link == "InfiniBand" and "ACTIVE" in state and rate and rate[0].isdigit(): + rates[device] = int(rate[0]) + fastest = max(rates.values(), default=None) + return [device for device, rate in rates.items() if rate == fastest] + + +def _read_attribute(directory: str, name: str) -> str: + try: + with open(os.path.join(directory, name)) as file: + return file.read().strip() + except OSError: + return "" + + +def _local_address(master: str) -> str: + """The address the route to the master leaves from. Not ``gethostbyname(gethostname())``, + which answers 127.0.1.1 where /etc/hosts maps the host's name to loopback.""" + for endpoint in (_split_address(master), _ROUTE_PROBE_ENDPOINT): + if endpoint is None: + continue + try: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as probe: + probe.connect(endpoint) + address = probe.getsockname()[0] + except OSError: + continue + if not address.startswith("127."): + return address + try: + return socket.gethostbyname(socket.gethostname()) + except OSError: + return "127.0.0.1" + + +def _wait_for_master(address: str) -> None: + """Turns an unreachable master into one line naming it, instead of a bare code from setup.""" + endpoint = _split_address(address) + if endpoint is None: + logger.warning(f"mooncake: cannot check master_server_address {address!r}: not host:port") + return + deadline = time.monotonic() + MASTER_WAIT_S + while True: + try: + with socket.create_connection(endpoint, timeout=1.0): + return + except OSError as error: + if time.monotonic() >= deadline: + raise TimeoutError( + f"the Mooncake master at {address} did not accept connections within " + f"{MASTER_WAIT_S:g} s: {error}" + ) from error + time.sleep(MASTER_RETRY_S) + + +def _split_address(address: str) -> tuple[str, int] | None: + host, _, port = address.rpartition(":") + return (host.strip("[]"), int(port)) if host and port.isdigit() else None diff --git a/tests/unittest/_torch/disaggregation/blob/test_mooncake_setup.py b/tests/unittest/_torch/disaggregation/blob/test_mooncake_setup.py new file mode 100644 index 000000000000..82f8b508680e --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_mooncake_setup.py @@ -0,0 +1,572 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Store config, opening the Mooncake SDK and driver, and the blob package's layering. + +The SDK is a fake module put into ``sys.modules``, the master probe and the route probe are +replaced, and the InfiniBand sysfs tree is a temporary directory, so nothing here needs a network, +a Mooncake install or a GPU. +""" + +from __future__ import annotations + +import ast +import pathlib +import socket +import sys +import threading +import types + +import pytest +import yaml +from fake_mooncake import FakeMooncakeStore, FakePool + +import tensorrt_llm._torch.disaggregation.blob as blob_package +from tensorrt_llm._torch.disaggregation.blob import mooncake_setup +from tensorrt_llm._torch.disaggregation.blob.backend import BlobConfig +from tensorrt_llm._torch.disaggregation.blob.mooncake import MooncakeDriver +from tensorrt_llm._torch.disaggregation.blob.mooncake_setup import ( + MooncakeConfig, + load_store_config, + open_mooncake_driver, + open_sdk, +) +from tensorrt_llm._torch.disaggregation.blob.store import RegionLayout +from tensorrt_llm._torch.disaggregation.testing.host_memory import HostBuffer + +# An exception that escapes a thread fails the test. +pytestmark = [ + pytest.mark.cpu_only, + pytest.mark.filterwarnings("error::pytest.PytestUnhandledThreadExceptionWarning"), +] + +MASTER = "10.0.0.1:50051" +DRAIN_S = 5.0 # a drained driver's worker threads have exited +ROUTED_ADDRESS = "10.0.0.7" + + +def store_config(protocol: str = "tcp", **mooncake) -> dict: + return { + "blob": {"key_prefix": "t"}, + "mooncake": { + "master_server_address": MASTER, + "global_segment_size": 0, + "protocol": protocol, + **mooncake, + }, + } + + +class Sdk: + """The fake ``mooncake.store`` module's state: every store it built, and the masters probed.""" + + def __init__(self) -> None: + self.pool = FakePool() + self.stores: list[FakeMooncakeStore] = [] + self.masters: list[str] = [] + self.setup_answer = 0 + + def build(self) -> FakeMooncakeStore: + store = FakeMooncakeStore(self.pool) + store.setup_answer = self.setup_answer + self.stores.append(store) + return store + + +def read_blob(section) -> BlobConfig: + return BlobConfig(**section) + + +def built_by(monkeypatch: pytest.MonkeyPatch, *builds: tuple[str, str]) -> None: + """Makes ``builds`` (name, version) the distributions that list the imported SDK module.""" + monkeypatch.setattr(mooncake_setup, "_builds_listing_file", lambda module: list(builds)) + + +@pytest.fixture +def sdk(monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> Sdk: + sdk = Sdk() + package = types.ModuleType("mooncake") + package.__path__ = [] + module = types.ModuleType("mooncake.store") + module.MooncakeDistributedStore = sdk.build + package.store = module + monkeypatch.setitem(sys.modules, "mooncake", package) + monkeypatch.setitem(sys.modules, "mooncake.store", module) + built_by(monkeypatch, ("mooncake-transfer-engine-cuda13", "0.3.13")) + monkeypatch.setattr(mooncake_setup, "_wait_for_master", sdk.masters.append) + monkeypatch.setattr(mooncake_setup, "_local_address", lambda master: ROUTED_ADDRESS) + monkeypatch.setattr(mooncake_setup, "IB_SYSFS_ROOT", str(tmp_path / "no-infiniband")) + for name in ("MC_TCP_PROTO", "MC_STORE_LOCAL_HOT_CACHE_SIZE", "MC_USE_TENT", "MC_USE_TEV1"): + monkeypatch.delenv(name, raising=False) + return sdk + + +def ib_device(root: pathlib.Path, name: str, link_layer: str, state: str, rate: str) -> None: + port = root / name / "ports" / "1" + port.mkdir(parents=True) + (port / "link_layer").write_text(f"{link_layer}\n") + (port / "state").write_text(f"{state}\n") + (port / "rate").write_text(f"{rate}\n") + + +# Store config. + + +def test_store_config_refuses_missing_and_unknown_sections_and_keys() -> None: + config = store_config() + del config["mooncake"]["global_segment_size"] + with pytest.raises(ValueError, match="'mooncake'.*global_segment_size"): + load_store_config(config) + with pytest.raises(ValueError, match="unknown sections"): + load_store_config({**store_config(), "mooncake_store": {}}) + with pytest.raises(ValueError, match="'mooncake' is a required mapping"): + load_store_config({"blob": {}}) + + +BAD_VALUES = { + "blob.key_prefx=t": ("blob", "key_prefx", "t"), # an unknown key + # How Mooncake's work runs is set in its own section. + "blob.workers=2": ("blob", "workers", 2), + "mooncake.devce_name=mlx5_0": ("mooncake", "devce_name", "mlx5_0"), + "mooncake.protocol=rdmaa": ("mooncake", "protocol", "rdmaa"), + "mooncake.global_segment_size=True": ("mooncake", "global_segment_size", True), + "mooncake.rows_per_call=64.0": ("mooncake", "rows_per_call", 64.0), + "mooncake.bytes_per_call=1.28e8": ("mooncake", "bytes_per_call", 1.28e8), + "blob.key_prefix=7": ("blob", "key_prefix", 7), + "mooncake.workers=2.5": ("mooncake", "workers", 2.5), + "mooncake.workers=True": ("mooncake", "workers", True), + "mooncake.workers=0": ("mooncake", "workers", 0), + "mooncake.global_segment_size=1.0": ("mooncake", "global_segment_size", 1.0), + "mooncake.local_buffer_size=-1": ("mooncake", "local_buffer_size", -1), + "mooncake.master_server_address=5": ("mooncake", "master_server_address", 5), + "mooncake.device_name=None": ("mooncake", "device_name", None), + # One master's host:port: an HA master list would leave nothing to wait for. + "mooncake.master_server_address=HA": ( + "mooncake", + "master_server_address", + "etcd://10.0.0.1:2379;10.0.0.2:2379", + ), +} + + +@pytest.mark.parametrize("case", BAD_VALUES) +def test_unknown_keys_and_values_of_the_wrong_type_are_rejected(sdk: Sdk, case: str) -> None: + section, name, value = BAD_VALUES[case] + config = store_config() + config[section][name] = value + with pytest.raises(ValueError, match=f"section '{section}'.*{name}"): + open_mooncake_driver(config, read_blob) + assert sdk.stores == [] and sdk.masters == [] # refused before the SDK was opened + + +def test_an_empty_blob_section_is_empty_and_anything_but_a_mapping_is_refused() -> None: + mooncake = store_config()["mooncake"] + for empty in ({"blob": None}, {}): # "blob:" alone, or no blob section + assert load_store_config({**empty, "mooncake": mooncake})[0] == {} + for not_a_mapping in ([], False, 0): + with pytest.raises(ValueError, match="'blob' an optional one"): + load_store_config({"blob": not_a_mapping, "mooncake": mooncake}) + with pytest.raises(ValueError, match="unknown sections"): # keys of mixed types + load_store_config({**store_config(), 1: "x", "extra": "y"}) + + +def test_an_unreadable_store_config_file_is_a_value_error(tmp_path: pathlib.Path) -> None: + with pytest.raises(ValueError, match="missing.yaml"): + load_store_config(tmp_path / "missing.yaml") + broken = tmp_path / "broken.yaml" + broken.write_text("mooncake: [") + with pytest.raises(ValueError, match="broken.yaml"): + load_store_config(broken) + + +def test_store_config_reads_a_yaml_file(tmp_path: pathlib.Path) -> None: + path = tmp_path / "store.yaml" + path.write_text(yaml.safe_dump(store_config(local_buffer_size=0, workers=2))) + blob, mooncake = load_store_config(path) + assert blob == {"key_prefix": "t"} + assert mooncake == MooncakeConfig(MASTER, 0, protocol="tcp", local_buffer_size=0, workers=2) + + +# Opening the SDK. + + +def test_open_passes_setup_arguments_by_name( + sdk: Sdk, tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch +) -> None: + sysfs = tmp_path / "infiniband" + ib_device(sysfs, "mlx5_0", "InfiniBand", "4: ACTIVE", "400 Gb/sec (4X NDR)") + ib_device(sysfs, "mlx5_1", "InfiniBand", "4: ACTIVE", "400 Gb/sec (4X NDR)") + monkeypatch.setattr(mooncake_setup, "IB_SYSFS_ROOT", str(sysfs)) + tunings = [] + + def driver_recording_its_tuning(store, **tuning) -> MooncakeDriver: + tunings.append(tuning) + return MooncakeDriver(store, **tuning) + + monkeypatch.setattr(mooncake_setup, "MooncakeDriver", driver_recording_its_tuning) + driver, blob = open_mooncake_driver(store_config("rdma", rows_per_call=4), read_blob) + (store,) = sdk.stores + (setup,) = store.calls_to("setup") + assert setup.args == () + assert setup.kwargs == { + "local_hostname": ROUTED_ADDRESS, # where the route to the master leaves from + "metadata_server": "P2PHANDSHAKE", + "global_segment_size": 0, + "local_buffer_size": 16 << 20, + "protocol": "rdma", + "rdma_devices": "mlx5_0,mlx5_1", # the fastest active InfiniBand ports + "master_server_addr": MASTER, + } + assert sdk.masters == [MASTER] + assert isinstance(driver, MooncakeDriver) + # The config's own values win over the driver's defaults. + assert tunings == [{"workers": 4, "rows_per_call": 4, "bytes_per_call": 128 << 20}] + assert blob == BlobConfig(key_prefix="t") + memory = HostBuffer(64) # the driver works through the client it was given + handle = driver.register(memory.address, memory.nbytes, RegionLayout(8, 8)) + assert store.calls_to("register_buffer") + driver.unregister(handle) + assert driver.drain(DRAIN_S) # its worker threads exit + + open_sdk(MooncakeConfig(MASTER, 1 << 30, "tcp", "eth9", "10.9.9.9", "etcd://m:2379", 0)) + assert sdk.stores[1].calls_to("setup")[0].kwargs == { + "local_hostname": "10.9.9.9", + "metadata_server": "etcd://m:2379", + "global_segment_size": 1 << 30, + "local_buffer_size": 0, + "protocol": "tcp", + "rdma_devices": "eth9", + "master_server_addr": MASTER, + } + + +def test_device_detection_picks_the_fastest_active_ib_ports(tmp_path: pathlib.Path) -> None: + ib_device(tmp_path, "mlx5_0", "InfiniBand", "4: ACTIVE", "400 Gb/sec (4X NDR)") + ib_device(tmp_path, "mlx5_1", "InfiniBand", "4: ACTIVE", "200 Gb/sec (4X HDR)") + ib_device(tmp_path, "mlx5_2", "Ethernet", "4: ACTIVE", "400 Gb/sec (4X NDR)") + ib_device(tmp_path, "mlx5_3", "InfiniBand", "1: DOWN", "400 Gb/sec (4X NDR)") + ib_device(tmp_path, "mlx5_4", "InfiniBand", "4: ACTIVE", "400 Gb/sec (4X NDR)") + assert mooncake_setup._fastest_active_ib_devices(str(tmp_path)) == ["mlx5_0", "mlx5_4"] + assert mooncake_setup._fastest_active_ib_devices(str(tmp_path / "missing")) == [] + + +def test_a_configured_device_is_used_as_given( + tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch +) -> None: + ib_device(tmp_path, "mlx5_0", "InfiniBand", "4: ACTIVE", "400 Gb/sec (4X NDR)") + monkeypatch.setattr(mooncake_setup, "IB_SYSFS_ROOT", str(tmp_path)) + configured = MooncakeConfig(MASTER, 0, protocol="rdma", device_name="mlx5_9") + assert mooncake_setup._rdma_devices(configured) == "mlx5_9" # not the detected mlx5_0 + + +def test_missing_sdk_reports_install_and_shadow_hint(sdk: Sdk, monkeypatch) -> None: + shadow = types.ModuleType("mooncake") + shadow.__file__ = "/opt/nixl/lib/python3/site-packages/mooncake/__init__.py" + monkeypatch.setitem(sys.modules, "mooncake", shadow) + monkeypatch.setitem(sys.modules, "mooncake.store", None) + with pytest.raises(ImportError) as missing: + open_sdk(MooncakeConfig(MASTER, 0, "tcp")) + message = str(missing.value) + assert "mooncake-transfer-engine-cuda13==0.3.13" in message + assert "/opt/nixl/" in message and "shadows" in message + with pytest.raises(ImportError, match="mooncake-transfer-engine-cuda13"): + open_mooncake_driver(store_config(), read_blob) + + +def test_tcp_legacy_framing_is_refused(sdk: Sdk, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("MC_TCP_PROTO", "1") + with pytest.raises(ValueError, match="MC_TCP_PROTO"): + open_mooncake_driver(store_config("tcp"), read_blob) + assert sdk.stores == [] + open_sdk(MooncakeConfig(MASTER, 0, "rdma", device_name="mlx5_0")) # RDMA does not frame + assert len(sdk.stores) == 1 + + +# The SDK reads the size with stoull, which takes leading spaces and stops at the first non-digit. +@pytest.mark.parametrize("size", ["64", " 64", "64MB", "1"]) +def test_local_hot_cache_is_refused(sdk: Sdk, monkeypatch: pytest.MonkeyPatch, size: str) -> None: + monkeypatch.setenv("MC_STORE_LOCAL_HOT_CACHE_SIZE", size) + with pytest.raises(ValueError, match="MC_STORE_LOCAL_HOT_CACHE_SIZE"): + open_mooncake_driver(store_config(), read_blob) + assert sdk.stores == [] + + +@pytest.mark.parametrize("size", ["", "0", " 0"]) +def test_local_hot_cache_off_is_accepted( + sdk: Sdk, monkeypatch: pytest.MonkeyPatch, size: str +) -> None: + monkeypatch.setenv("MC_STORE_LOCAL_HOT_CACHE_SIZE", size) + open_sdk(MooncakeConfig(MASTER, 0, "tcp")) + assert len(sdk.stores) == 1 + + +@pytest.mark.parametrize("variable", ["MC_USE_TENT", "MC_USE_TEV1"]) +def test_the_newer_transfer_engine_is_refused( + sdk: Sdk, monkeypatch: pytest.MonkeyPatch, variable: str +) -> None: + monkeypatch.setenv(variable, "") # the SDK only asks whether it is set + with pytest.raises(ValueError, match=f"{variable} selects a transfer engine"): + open_mooncake_driver(store_config(), read_blob) + assert sdk.stores == [] + + +CUDA13 = "mooncake-transfer-engine-cuda13" +SDK_BUILDS = { + # name: (the distributions that list the imported SDK module, the one refused or None) + "cuda13": ([(CUDA13, "0.3.13")], None), + "cuda12": ([("mooncake-transfer-engine", "0.3.13")], None), + "non_cuda": ([("mooncake-transfer-engine-non-cuda", "0.3.13")], None), + "cuda12_beside_cuda13": ([(CUDA13, "0.3.13"), ("mooncake-transfer-engine", "0.3.13")], None), + "newer": ([(CUDA13, "0.3.14")], f"{CUDA13}==0.3.14"), + "efa_transport": ([("mooncake-transfer-engine-efa", "0.3.13")], "efa==0.3.13"), + # Installed after the pinned build, the other one's file is the module on disk. + "older_beside_pinned": ( + [(CUDA13, "0.3.13"), ("mooncake-transfer-engine", "0.3.11")], + "mooncake-transfer-engine==0.3.11", + ), + "no_distribution": ([], "no installed distribution"), # a source tree or a shadowing package +} + + +@pytest.mark.parametrize("case", SDK_BUILDS) +def test_only_pinned_sdk_builds_are_opened( + sdk: Sdk, monkeypatch: pytest.MonkeyPatch, case: str +) -> None: + builds, refused = SDK_BUILDS[case] + built_by(monkeypatch, *builds) + if refused is None: + driver, _ = open_mooncake_driver(store_config(), read_blob) + assert len(sdk.stores) == 1 + assert driver.drain(DRAIN_S) # its worker threads exit + return + with pytest.raises(ImportError, match=f"comes from .*{refused}.*follows 0.3.13 of"): + open_mooncake_driver(store_config(), read_blob) + assert sdk.stores == [] and sdk.masters == [] + + +def install(site: pathlib.Path, dist_info: str, metadata: str, files: list[str]) -> None: + """Writes an installed distribution's metadata into ``site``, listing ``files``.""" + info = site / f"{dist_info}.dist-info" + info.mkdir(parents=True) + (info / "METADATA").write_text(metadata) + (info / "RECORD").write_text("".join(f"{file},,\n" for file in files)) + + +def test_every_build_listing_the_imported_module_is_found( + tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch +) -> None: + real = tmp_path / "real" + (real / "mooncake").mkdir(parents=True) + (real / "mooncake" / "store.so").touch() + store = "mooncake/store.so" + install(real, "nixl-0.9", "Name: nixl\nVersion: 0.9\n", ["mooncake/__init__.py"]) # a shadow + install( + real, "M_T_E-0.3.13", "Name: Mooncake_Transfer_Engine_CUDA13\nVersion: 0.3.13\n", [store] + ) + # A reinstall over another version can leave a second dist-info of the same name. + install(real, "m_t_e-0.3.12", f"Name: {CUDA13}\nVersion: 0.3.12\n", [store]) + install(real, "unnamed-1.0", "Version: 1.0\n", [store]) + site = tmp_path / "site" + site.symlink_to(real) # the distributions are found under a symlinked site directory + monkeypatch.setattr(sys, "path", [str(site)]) + module = types.ModuleType("mooncake.store") + module.__file__ = str(real / "mooncake" / "store.so") + assert sorted(mooncake_setup._builds_listing_file(module)) == [ + ("a distribution without a name", "1.0"), + (CUDA13, "0.3.12"), + (CUDA13, "0.3.13"), + ] + with pytest.raises(ImportError, match=f"{CUDA13}==0.3.12"): + mooncake_setup._require_pinned_builds(module) + module.__file__ = str(real / "elsewhere" / "store.so") + assert mooncake_setup._builds_listing_file(module) == [] + del module.__file__ + assert mooncake_setup._builds_listing_file(module) == [] + + +class RoutedProbe: + """A datagram socket whose connect to an endpoint in ``routes`` leaves from the address given + there; any other endpoint has no route.""" + + routes: dict[tuple[str, int], str] = {} + + def __init__(self, *args) -> None: + self.address = "" + + def __enter__(self) -> RoutedProbe: + return self + + def __exit__(self, *exc) -> None: + return None + + def connect(self, endpoint: tuple[str, int]) -> None: + if endpoint not in self.routes: + raise OSError("no route") + self.address = self.routes[endpoint] + + def getsockname(self) -> tuple[str, int]: + return self.address, 0 + + +ROUTE_PROBE = mooncake_setup._ROUTE_PROBE_ENDPOINT +MASTER_ENDPOINT = ("10.0.0.1", 50051) +LOCAL_ADDRESS_CASES = { + # name: (the routes off this host, the address this host's name resolves to, expected) + "route_to_the_master": ({MASTER_ENDPOINT: "10.0.0.7"}, None, "10.0.0.7"), + "loopback_to_the_master": ( + {MASTER_ENDPOINT: "127.0.0.1", ROUTE_PROBE: "10.0.0.8"}, + None, + "10.0.0.8", + ), + "no_route": ({}, "10.0.0.9", "10.0.0.9"), + "nothing_resolves": ({}, None, "127.0.0.1"), +} + + +@pytest.mark.parametrize("case", LOCAL_ADDRESS_CASES) +def test_the_local_address_is_where_the_route_to_the_master_leaves( + monkeypatch: pytest.MonkeyPatch, case: str +) -> None: + routes, resolved, expected = LOCAL_ADDRESS_CASES[case] + + def gethostbyname(name: str) -> str: + if resolved is None: + raise socket.gaierror("unknown host") + return resolved + + probe = type("Probe", (RoutedProbe,), {"routes": routes}) + fake = types.SimpleNamespace( + socket=probe, + AF_INET=socket.AF_INET, + SOCK_DGRAM=socket.SOCK_DGRAM, + gethostbyname=gethostbyname, + gethostname=lambda: "this-host", + ) + monkeypatch.setattr(mooncake_setup, "socket", fake) + assert mooncake_setup._local_address(MASTER) == expected + + +def test_waiting_for_the_master_ends_when_it_accepts_and_names_it_when_it_never_does( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with socket.create_server(("127.0.0.1", 0)) as server: + address = "127.0.0.1:%d" % server.getsockname()[1] + mooncake_setup._wait_for_master(address) # accepts at once + monkeypatch.setattr(mooncake_setup, "MASTER_WAIT_S", 0.2) + monkeypatch.setattr(mooncake_setup, "MASTER_RETRY_S", 0.05) + with pytest.raises(TimeoutError, match=f"master at {address} did not accept"): + mooncake_setup._wait_for_master(address) # the port is closed now + + +def test_waiting_for_the_master_outlasts_one_that_comes_up_late( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with socket.create_server(("127.0.0.1", 0)) as probe: + port = probe.getsockname()[1] # free once the probe closes: nothing listens yet + servers: list[socket.socket] = [] + master_starts = threading.Timer( + 0.3, lambda: servers.append(socket.create_server(("127.0.0.1", port))) + ) + monkeypatch.setattr(mooncake_setup, "MASTER_WAIT_S", 5.0) + monkeypatch.setattr(mooncake_setup, "MASTER_RETRY_S", 0.05) + master_starts.start() + try: + mooncake_setup._wait_for_master(f"127.0.0.1:{port}") + finally: + master_starts.join() + for server in servers: + server.close() + + +def test_setup_failure_closes_the_half_open_client(sdk: Sdk) -> None: + sdk.setup_answer = -1 + with pytest.raises(RuntimeError, match="setup failed with -1"): + open_mooncake_driver(store_config(), read_blob) + (store,) = sdk.stores + assert len(store.calls_to("close")) == 1 + + +def test_a_driver_that_cannot_start_closes_the_client_it_was_given( + sdk: Sdk, monkeypatch: pytest.MonkeyPatch +) -> None: + start = threading.Thread.start + + def refuse_worker_threads(thread: threading.Thread) -> None: + if thread.name.startswith("mooncake-driver-"): + raise RuntimeError("can't start new thread") + start(thread) + + monkeypatch.setattr(threading.Thread, "start", refuse_worker_threads) + with pytest.raises(RuntimeError, match="can't start new thread"): + open_mooncake_driver(store_config(), read_blob) + (store,) = sdk.stores + assert store.calls_to("setup") and len(store.calls_to("close")) == 1 + + +# Layering. + + +BLOB_PACKAGE = blob_package.__name__ +DISAGGREGATION = BLOB_PACKAGE.rsplit(".", 1)[0] +THIRD_PARTY_ALLOWED = {"numpy", "yaml"} + + +def _imports(path: pathlib.Path) -> list[tuple[str, str | None]]: + """Every module ``path`` imports, absolute, with the function it is imported in (or None).""" + found: list[tuple[str, str | None]] = [] + + class Visitor(ast.NodeVisitor): + function: str | None = None + + def visit_FunctionDef(self, node: ast.FunctionDef) -> None: + outer, self.function = self.function, node.name + self.generic_visit(node) + self.function = outer + + def visit_Import(self, node: ast.Import) -> None: + found.extend((alias.name, self.function) for alias in node.names) + + def visit_ImportFrom(self, node: ast.ImportFrom) -> None: + module = node.module or "" + if node.level: # relative to the blob package + anchor = BLOB_PACKAGE.rsplit(".", node.level - 1)[0] + if not module: # "from . import a, b" imports the modules a and b + found.extend((f"{anchor}.{alias.name}", self.function) for alias in node.names) + return + module = f"{anchor}.{module}" + found.append((module, self.function)) + + Visitor().visit(ast.parse(path.read_text())) + return found + + +def test_blob_modules_import_only_their_layer() -> None: + root = pathlib.Path(blob_package.__file__).parent + modules = {path.stem: _imports(path) for path in sorted(root.glob("*.py"))} + assert {"backend", "mooncake", "mooncake_setup", "store", "keys"} <= set(modules) + sdk_imports = [] + for name, imports in modules.items(): + for module, function in imports: + top = module.split(".")[0] + where = f"blob/{name}.py imports {module}" + if top == "mooncake": + sdk_imports.append((name, function)) + continue + assert top != "torch" and not module.startswith("tensorrt_llm._torch.pyexecutor"), where + if top == "tensorrt_llm": + assert module == "tensorrt_llm.logger" or module.startswith(DISAGGREGATION), where + else: + assert top in sys.stdlib_module_names or top in THIRD_PARTY_ALLOWED, where + assert sdk_imports == [("mooncake_setup", "open_sdk")] + + def package_imports(name: str) -> set[str]: + return {m for m, _ in modules[name] if m.startswith(BLOB_PACKAGE + ".")} + + assert package_imports("backend") == {f"{BLOB_PACKAGE}.keys", f"{BLOB_PACKAGE}.store"} + assert package_imports("mooncake") == {f"{BLOB_PACKAGE}.store"} + assert package_imports("mooncake_setup") == {f"{BLOB_PACKAGE}.mooncake"} + memory_kind = f"{DISAGGREGATION}.base.memory" # all the boundary shares with the contract + for name in ("store", "keys"): + for module, _ in modules[name]: + top = module.split(".")[0] + columns = name == "store" and top == "numpy" # a RowBatch holds its rows as columns + assert top in sys.stdlib_module_names or columns or module == memory_kind, name From ec265c74665f9c5ef8a1fb1f43fa05b48be62cb0 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Tue, 27 Oct 2026 08:23:13 +0000 Subject: [PATCH 03/10] [None][test] Run the blob backend over a real Mooncake pool The fake SDK shows what the driver asks for; only a real master shows whether the SDK answers that way. test_mooncake_real.py runs the shared driver_conformance checks over a real pool, one fresh client per check, with the scenarios a real master can stage (a key another client is still writing, a get or put held inside the SDK), and runs the backend between real clients for what the fake models beyond them: a full pool that declines rows, a get that lands a shorter object's bytes and refuses a longer one, a second publish of keys whose sessions the same client still holds, also once the first put is stuck past the master's put_start_discard_timeout_sec, and four workers moving 2000 rows byte for byte. The clients give the pool no segment and keep no local buffer, so every object sits in another process's segment (running_segment_provider in test_common/mooncake_utils.py) and nothing falls back to the SDK's copying path. The file runs only with TRTLLM_MOONCAKE_REAL=1 and fails unless mooncake_master and the client are both a pinned 0.3.13 build; it records both versions as test-suite properties. An exception escaping a thread fails the test. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../_torch/disaggregation/blob/README.md | 4 +- tests/test_common/mooncake_utils.py | 80 +++- .../disaggregation/blob/test_mooncake_real.py | 449 ++++++++++++++++++ 3 files changed, 530 insertions(+), 3 deletions(-) create mode 100644 tests/unittest/_torch/disaggregation/blob/test_mooncake_real.py diff --git a/tensorrt_llm/_torch/disaggregation/blob/README.md b/tensorrt_llm/_torch/disaggregation/blob/README.md index ae06c7aba0d0..975f319a731d 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/README.md +++ b/tensorrt_llm/_torch/disaggregation/blob/README.md @@ -42,7 +42,7 @@ store SDK pytest tests/unittest/_torch/disaggregation/blob ``` -All of it is CPU-only. The files, top down: +All of it is CPU-only; `test_mooncake_real.py` also needs a Mooncake pool. The files, top down: - `test_blob_example.py`: the backend end to end, each test the smallest example of one part of the picture: two endpoints sharing objects, a lease holder's decisions and leases, the @@ -58,6 +58,8 @@ All of it is CPU-only. The files, top down: answers as 0.3.13 does, the conformance suite included. - `test_mooncake_setup.py`: the store config, the SDK's setup and the refused settings and builds, and that each module imports only its layer. +- `test_mooncake_real.py`: the Mooncake driver and the backend over a real pool; runs only with + `TRTLLM_MOONCAKE_REAL=1`. ## Writing a driver diff --git a/tests/test_common/mooncake_utils.py b/tests/test_common/mooncake_utils.py index e50af9fa7ba5..1b70315724dd 100644 --- a/tests/test_common/mooncake_utils.py +++ b/tests/test_common/mooncake_utils.py @@ -12,7 +12,7 @@ # 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. -"""Running a Mooncake master for a test, on ports picked for it. +"""Running a Mooncake master for a test, on ports picked for it, and the memory behind it. A deployment names its master's ports, so `running_master` takes them as given. A test has no pair to name and wants any free one, on a machine whose @@ -20,20 +20,44 @@ binding it and letting go, so another process can still take it before the master binds, and the master then exits before it accepts connections. Each attempt picks a fresh pair. + +The master stores nothing itself: `running_segment_provider` gives a pool its +memory from a separate process. """ import contextlib +import os +import subprocess # nosec B404 +import sys +import time from typing import Iterator +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import local_address from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.master import ( LaunchedMaster, + resolve_device_name, running_master, ) from tensorrt_llm._utils import get_free_ports -__all__ = ["running_master_on_free_ports"] +__all__ = ["running_master_on_free_ports", "running_segment_provider"] PORT_ATTEMPTS = 5 +SEGMENT_PROVIDER_TIMEOUT = 60.0 + +# Runs in the provider process: join the pool, mount the segment, and hold it until stdin closes. +_SEGMENT_PROVIDER = """ +import pathlib, sys +from mooncake.store import MooncakeDistributedStore +master, segment_size, protocol, devices, hostname, ready = sys.argv[1:] +store = MooncakeDistributedStore() +code = store.setup(hostname, "P2PHANDSHAKE", int(segment_size), 16 << 20, protocol, devices, master) +if code != 0: + sys.exit(f"segment provider: setup failed with {code}") +pathlib.Path(ready).touch() +sys.stdin.read() +store.close() +""" def _start_master( @@ -75,3 +99,55 @@ def running_master_on_free_ports( """ with contextlib.ExitStack() as stack: yield _start_master(stack, run_dir, attempts, kwargs) + + +@contextlib.contextmanager +def running_segment_provider( + master_address: str, + run_dir: str, + *, + segment_size: int, + protocol: str = "tcp", + device_name: str = "", + timeout: float = SEGMENT_PROVIDER_TIMEOUT, +) -> Iterator[subprocess.Popen]: + """Give the pool a `segment_size`-byte segment from a separate process, for the body's duration. + + A client with no segment of its own then finds every object in another + process's memory, so each transfer goes through the transfer engine as it + would across hosts, never through a copy inside one process. + + Args: + master_address: The pool's master, `host:port`. + run_dir: Where the provider's log and readiness marker go. + segment_size: Bytes of the provider's segment. + protocol: `rdma` or `tcp`. + device_name: RDMA devices; empty picks the host's fastest InfiniBand ones. + timeout: Seconds to wait for the provider to join the pool. + """ + os.makedirs(run_dir, exist_ok=True) + ready = os.path.join(run_dir, "segment_provider.ready") + log_path = os.path.join(run_dir, "segment_provider.log") + devices = resolve_device_name(protocol, device_name) + command = [sys.executable, "-c", _SEGMENT_PROVIDER, master_address, str(segment_size)] + command += [protocol, devices, local_address(master_address), ready] + with open(log_path, "wb") as log: + process = subprocess.Popen( # nosec B603 + command, stdin=subprocess.PIPE, stdout=log, stderr=subprocess.STDOUT + ) + try: + deadline = time.monotonic() + timeout + while not os.path.exists(ready): + if process.poll() is not None or time.monotonic() > deadline: + with open(log_path, errors="replace") as log: + tail = log.read()[-4000:] + raise RuntimeError(f"the segment provider did not join {master_address}:\n{tail}") + time.sleep(0.1) + yield process + finally: + process.stdin.close() + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + process.kill() + process.wait() diff --git a/tests/unittest/_torch/disaggregation/blob/test_mooncake_real.py b/tests/unittest/_torch/disaggregation/blob/test_mooncake_real.py new file mode 100644 index 000000000000..840480e8bae9 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_mooncake_real.py @@ -0,0 +1,449 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The blob backend over a real Mooncake pool: a master, and a separate process holding objects. + +The tested clients give the pool no segment (``global_segment_size`` 0), so every object sits in +the other process's segment and every transfer goes through the transfer engine; their +``local_buffer_size`` is 0, so a call that needed the SDK's copying path would fail. Runs only with +``TRTLLM_MOONCAKE_REAL=1``, on a host with ``mooncake_master`` and ``mooncake.store``. The +transport is ``TRTLLM_MOONCAKE_PROTOCOL`` (default rdma) over ``TRTLLM_MOONCAKE_DEVICES`` (default: +the host's fastest active InfiniBand devices). The master must be shipped by one of the pinned SDK +distributions and report 0.3.13, and every installed pinned client build must be 0.3.13; both +versions are printed and recorded as test-suite properties. The shared +``driver_conformance`` checks run here too, each over a fresh client and under a key prefix of its +own. +""" + +from __future__ import annotations + +import contextlib +import importlib.metadata +import importlib.util +import os +import pathlib +import re +import secrets +import shlex +import shutil +import subprocess # nosec B404 +import threading +import time +from collections.abc import Sequence + +import numpy as np +import pytest +from cpu_fixtures import WAIT_S, Endpoint, masks, store_key +from fake_mooncake import KeyGatedClient + +from tensorrt_llm._torch.disaggregation.base.backend import Outcome +from tensorrt_llm._torch.disaggregation.blob.mooncake import MooncakeDriver +from tensorrt_llm._torch.disaggregation.blob.mooncake_setup import ( + PINNED_SDK_DISTRIBUTIONS, + PINNED_SDK_VERSION, + MooncakeConfig, + open_sdk, +) +from tensorrt_llm._torch.disaggregation.testing.driver_conformance import SLOT as CHECK_SLOT +from tensorrt_llm._torch.disaggregation.testing.driver_conformance import DriverConformance + +# An exception that escapes a thread fails the test. +pytestmark = [ + pytest.mark.skipif( + os.environ.get("TRTLLM_MOONCAKE_REAL") != "1", + reason="needs a real Mooncake pool; set TRTLLM_MOONCAKE_REAL=1", + ), + pytest.mark.filterwarnings("error::pytest.PytestUnhandledThreadExceptionWarning"), +] + +PROTOCOL = os.environ.get("TRTLLM_MOONCAKE_PROTOCOL", "rdma") +DEVICES = os.environ.get("TRTLLM_MOONCAKE_DEVICES", "") +POOL_BYTES = 1 << 30 +SLOT_BYTES = 64 << 10 + + +def _require_mooncake() -> str: + """Skips without a master or an SDK; fails unless both are the pinned build. Returns a line + naming both versions, for the run's log.""" + master = shutil.which("mooncake_master") + if master is None: + pytest.skip("mooncake_master is not on PATH") + try: + found = importlib.util.find_spec("mooncake.store") is not None + except ImportError: + found = False + if not found: + pytest.skip("mooncake.store does not import") + clients, owners = {}, [] + for name in PINNED_SDK_DISTRIBUTIONS: + try: + build = importlib.metadata.distribution(name) + except importlib.metadata.PackageNotFoundError: + continue + clients[name] = build.version + shipped = (pathlib.Path(build.locate_file(file)).resolve() for file in build.files or ()) + if pathlib.Path(master).resolve() in shipped: + owners.append(name) + answer = subprocess.run( # nosec B603 + [master, "--version"], capture_output=True, text=True, timeout=WAIT_S + ).stdout + reported = re.search(r"version (\S+)", answer) + versions = ( + f"mooncake_master {master} from {owners or 'no SDK build'}, reporting " + f"{reported.group(1) if reported else answer!r}; client builds {clients}" + ) + pinned = ( + set(clients.values()) == {PINNED_SDK_VERSION} + and reported + and reported[1] == PINNED_SDK_VERSION + ) + assert owners and pinned, f"not the pinned Mooncake {PINNED_SDK_VERSION}: {versions}" + return versions + + +def _master_binary(run_dir: str, flags: Sequence[str]) -> str: + """``mooncake_master``, or with ``flags`` a script in ``run_dir`` that starts it with them.""" + master = shutil.which("mooncake_master") + if not flags: + return master + script = pathlib.Path(run_dir) / "mooncake_master_with_flags" + script.write_text(f'#!/bin/sh\nexec {shlex.quote(master)} "$@" {shlex.join(flags)}\n') + script.chmod(0o755) + return str(script) + + +def _running_pool( + run_dir: str, segment_size: int, master_flags: Sequence[str] = () +) -> tuple[contextlib.ExitStack, str]: + """A master and a segment provider: the stack that stops them, and the master's address.""" + from test_common.mooncake_utils import running_master_on_free_ports, running_segment_provider + + stack = contextlib.ExitStack() + # Ratio 0: the master evicts only above its high watermark, never just because a put failed. + master = stack.enter_context( + running_master_on_free_ports( + run_dir, + protocol=PROTOCOL, + eviction_ratio=0.0, + binary=_master_binary(run_dir, master_flags), + ) + ) + stack.enter_context( + running_segment_provider( + master.address, + run_dir, + segment_size=segment_size, + protocol=PROTOCOL, + device_name=DEVICES, + ) + ) + return stack, master.address + + +@pytest.fixture(scope="module") +def pool(tmp_path_factory, record_testsuite_property): + versions = _require_mooncake() + print(versions) + record_testsuite_property("mooncake", versions) + stack, address = _running_pool(str(tmp_path_factory.mktemp("pool")), POOL_BYTES) + with stack: + yield address + + +class Clients: + """SDK clients of one master that give the pool no segment, and the endpoints built over them. + + ``close`` closes the clients only once every endpoint's shutdown returns no labels: a backend + thread still inside the SDK would crash on a closed client, so a failed test leaves them open. + """ + + def __init__(self, master: str) -> None: + self.master = master + self.opened = [] + self.endpoints: list[Endpoint] = [] + + def open_client(self): + self.opened.append(open_sdk(client_config(self.master))) + return self.opened[-1] + + def endpoint(self, names, *, sdk=None, workers: int = 2, **options) -> Endpoint: + """An endpoint over ``sdk``, a new client by default, with a driver of ``workers``.""" + options = {"slot_bytes": SLOT_BYTES, **options} + driver = MooncakeDriver(sdk or self.open_client(), workers=workers) + return self.add(Endpoint(driver, names, **options)) + + def add(self, endpoint: Endpoint) -> Endpoint: + self.endpoints.append(endpoint) + return endpoint + + def close(self) -> None: + if all(endpoint.backend.shutdown(WAIT_S) == frozenset() for endpoint in self.endpoints): + for store in self.opened: + store.close() + + +def client_config(master: str) -> MooncakeConfig: + """A client with no segment of its own and no buffer for the SDK's copying calls.""" + return MooncakeConfig(master, 0, PROTOCOL, DEVICES, local_buffer_size=0) + + +@pytest.fixture +def clients(pool): + opened = Clients(pool) + yield opened + opened.close() + + +def write_random_rows(endpoint: Endpoint, seed: int) -> dict[str, bytes]: + generator = np.random.default_rng(seed) + rows = {name: generator.bytes(endpoint.region.slot_bytes) for name in endpoint.names} + for name, data in rows.items(): + endpoint.write(name, data) + return rows + + +def finish(endpoint: Endpoint, submit, *names: str) -> Outcome: + attempt = submit(endpoint.view(*names), None) + outcome = endpoint.wait_outcome(attempt) + endpoint.wait_quiesced(attempt) + return outcome + + +def test_real_store_full_is_declined(tmp_path, request) -> None: + _require_mooncake() + # A pool of its own that one publish overfills: a 16 MiB segment (the smallest a TCP client + # may give, one allocator slab) holds at most 256 of the 512 rows. A PutStart that finds no room + # answers at once. The master evicts once the pool passes its high watermark, so how many rows + # land varies, but some are always declined and none fault. + stack, master = _running_pool(str(tmp_path), segment_size=16 << 20) + clients = Clients(master) + with stack: + names = [f"{request.node.name}-{i}" for i in range(512)] + writer = clients.endpoint(names, workers=1) + write_random_rows(writer, seed=4) + outcome = finish(writer, writer.backend.publish, *names) + delivered = int(outcome.delivered_masks[0].sum()) + declined = writer.driver.inner.stats()["rows_declined_pool_full"] + assert outcome.fault is None + assert delivered > 0 and declined > 0 and delivered + declined == len(names) + writer.close() + clients.close() # while the master still answers + + +def test_real_get_faults_a_shorter_object_it_landed_and_refuses_a_longer_one( + clients, request +) -> None: + # batch_get_into reads an object shorter than the buffer into the buffer's head and answers + # its size; it refuses a longer one before any transfer (INVALID_PARAMS). + short, long = f"{request.node.name}-short", f"{request.node.name}-long" + half = clients.endpoint([short], slot_bytes=SLOT_BYTES // 2) + double = clients.endpoint([long], slot_bytes=2 * SLOT_BYTES) + head = write_random_rows(half, seed=9)[short] + write_random_rows(double, seed=10) + assert masks(finish(half, half.backend.publish, short)) == [[True]] + assert masks(finish(double, double.backend.publish, long)) == [[True]] + reader = clients.endpoint([short, long], fill=0x2D) + outcome = finish(reader, reader.backend.fetch, short) # access ends: finish waits for it + assert masks(outcome) == [[False]] + assert f"object is {SLOT_BYTES // 2} bytes, expected {SLOT_BYTES}" in outcome.fault + assert reader.read(short) == head + b"-" * (SLOT_BYTES // 2) + outcome = finish(reader, reader.backend.fetch, long) + assert masks(outcome) == [[False]] and "INVALID_PARAMS(-600)" in outcome.fault + assert reader.read(long) == b"-" * SLOT_BYTES + for each in (half, double, reader): + each.close() + + +class TestRealMooncakeConformance(DriverConformance): + """The shared driver checks over the real pool, one fresh client per check, with the scenarios + a real master can stage: a key another client is writing, and a get or a put held inside the + SDK.""" + + timeout = 30.0 + + @pytest.fixture(autouse=True) + def _master(self, pool): + self.master = pool + + def make_driver(self) -> MooncakeDriver: + return MooncakeDriver(open_sdk(client_config(self.master))) + + def make_held_key(self): + key = f"held{secrets.token_hex(4)}/v1/H" # the pool outlives each check + writer = open_sdk(client_config(self.master)) # another client, still writing the key + assert writer.batch_put_session_start([key], [CHECK_SLOT]) == [0] + + def release() -> None: + try: + assert writer.batch_put_session_revoke([key]) == [0] + finally: + writer.close() + + return self.make_driver(), key, release + + def make_blocked_get(self): + key = f"blocked{secrets.token_hex(4)}/v1/S" + sdk = KeyGatedClient(open_sdk(client_config(self.master)), "batch_get_into", key) + return MooncakeDriver(sdk), key, sdk.entered, sdk.release.set + + def make_blocked_put(self): + key = f"blocked{secrets.token_hex(4)}/v1/P" + sdk = KeyGatedClient( + open_sdk(client_config(self.master)), "batch_put_from_multi_buffer_ranges", key + ) + return MooncakeDriver(sdk), key, sdk.entered, sdk.release.set + + +class DataCallsInFlight: + """An SDK client that counts its row-moving calls running at once; ``peak`` is the most.""" + + def __init__(self, store) -> None: + self.store = store + self.peak = self._running = 0 + self._lock = threading.Lock() + + def __getattr__(self, name: str): + return getattr(self.store, name) + + def batch_get_into(self, *args): + return self._counted(self.store.batch_get_into, args) + + def batch_put_from_multi_buffer_ranges(self, *args): + return self._counted(self.store.batch_put_from_multi_buffer_ranges, args) + + def _counted(self, call, args): + with self._lock: + self._running += 1 + self.peak = max(self.peak, self._running) + try: + return call(*args) + finally: + with self._lock: + self._running -= 1 + + +def test_real_four_workers_round_trip_2000_rows_without_crossed_bytes(clients, request) -> None: + names = [f"{request.node.name}-{i}" for i in range(2000)] + writer_sdk = DataCallsInFlight(clients.open_client()) + reader_sdk = DataCallsInFlight(clients.open_client()) + # Parts of 64 rows, so the one 2000-row operation runs on several workers at once. + writer = clients.add( + Endpoint( + MooncakeDriver(writer_sdk, workers=4, rows_per_call=64), names, slot_bytes=SLOT_BYTES + ) + ) + reader = clients.add( + Endpoint( + MooncakeDriver(reader_sdk, workers=4, rows_per_call=64), + names, + slot_bytes=SLOT_BYTES, + fill=0x2D, + ) + ) + rows = write_random_rows(writer, seed=6) + outcome = finish(writer, writer.backend.publish, *names) + assert outcome.fault is None and all(masks(outcome)[0]) + outcome = finish(reader, reader.backend.fetch, *names) + assert outcome.fault is None and all(masks(outcome)[0]) + crossed = [name for name in names if reader.read(name) != rows[name]] + assert crossed == [] + assert writer_sdk.peak >= 2 and reader_sdk.peak >= 2 + writer.close() + reader.close() + + +class FirstScatterWaits: + """An SDK client whose first put scatter waits, after its session started, for ``release`` + (at most ``hold_s``); ``starts`` lists the keys of every put session start.""" + + def __init__(self, store, hold_s: float = WAIT_S) -> None: + self.store = store + self.hold_s = hold_s + self.entered = threading.Event() + self.release = threading.Event() + self.starts: list[list[str]] = [] + self._first = threading.Lock() + + def __getattr__(self, name: str): + return getattr(self.store, name) + + def batch_put_session_start(self, keys, *args): + self.starts.append(list(keys)) + return self.store.batch_put_session_start(keys, *args) + + def batch_put_from_multi_buffer_ranges(self, *args): + if self._first.acquire(blocking=False): + self.entered.set() + self.release.wait(self.hold_s) + return self.store.batch_put_from_multi_buffer_ranges(*args) + + +def test_real_second_publish_of_keys_still_being_put_is_skipped_without_asking_the_sdk( + clients, request +) -> None: + # The first publish holds every key's put session while its scatter waits, so the driver + # skips the second publish's rows without asking the SDK; then the first publish delivers + # them all. + names = [f"{request.node.name}-{i}" for i in range(64)] + sdk = FirstScatterWaits(clients.open_client()) + # Three workers: puts leave one to gets, so the two publishes run side by side. + writer = clients.endpoint(names, sdk=sdk, workers=3) + rows = write_random_rows(writer, seed=8) + first = writer.backend.publish(writer.view(*names), None) + try: + assert sdk.entered.wait(WAIT_S) + second = writer.backend.publish(writer.view(*names), None) + skipped = writer.wait_outcome(second) + assert skipped.fault is None and masks(skipped) == [[False] * 64] + finally: + sdk.release.set() + assert masks(writer.wait_outcome(first)) == [[True] * 64] + for attempt in (first, second): + writer.wait_quiesced(attempt) + assert writer.driver.inner.stats()["rows_not_ready_held_by_own_put"] == 64 + assert len(sdk.starts) == 1 + reader = clients.endpoint(names, fill=0x2D) + assert all(masks(finish(reader, reader.backend.fetch, *names))[0]) + assert all(reader.read(name) == rows[name] for name in names) + writer.close() + reader.close() + + +DISCARD_S = 2 # the master's put_start_discard_timeout_sec in the test below; 30 s by default + + +def test_real_put_stuck_past_the_discard_timeout_keeps_its_keys(tmp_path, request) -> None: + # Past put_start_discard_timeout_sec the master would let another client's start discard a + # put still being written. This client's own second put of the keys is still skipped, and the + # objects never read as existing until the first put ends. + _require_mooncake() + flags = [f"--put_start_discard_timeout_sec={DISCARD_S}"] + stack, master = _running_pool(str(tmp_path), 64 << 20, master_flags=flags) + clients = Clients(master) + with stack: + names = [f"{request.node.name}-{i}" for i in range(8)] + keys = [store_key(name) for name in names] + sdk = FirstScatterWaits(clients.open_client(), hold_s=DISCARD_S + WAIT_S) + writer = clients.endpoint(names, sdk=sdk, workers=3) + observer = clients.open_client() + rows = write_random_rows(writer, seed=11) + first = writer.backend.publish(writer.view(*names), None) + try: + assert sdk.entered.wait(WAIT_S) + time.sleep(DISCARD_S + 1) + second = writer.backend.publish(writer.view(*names), None) + skipped = writer.wait_outcome(second) + assert skipped.fault is None and masks(skipped) == [[False] * 8] + assert observer.batch_is_exist(keys) == [0] * 8 + finally: + sdk.release.set() + assert masks(writer.wait_outcome(first)) == [[True] * 8] + for attempt in (first, second): + writer.wait_quiesced(attempt) + assert len(sdk.starts) == 1 # the second put never asked the SDK + assert observer.batch_is_exist(keys) == [1] * 8 + reader = clients.endpoint(names, fill=0x2D) + assert all(masks(finish(reader, reader.backend.fetch, *names))[0]) + assert all(reader.read(name) == rows[name] for name in names) + writer.close() + reader.close() + clients.close() # while the master still answers From 4d311477bb9f788e3fad498623bed6513390775b Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Tue, 27 Oct 2026 08:23:19 +0000 Subject: [PATCH 04/10] [None][test] Round-trip KV blocks between two staging lenders through the blob backend The CPU and real-pool tests drive the backend with host slots of their own; none shows that what the staging lender lends is what the backend moves. test_blob_backend_round_trip.py does it end to end on a GPU: one real KV cache manager publishes a request's blocks through its lender's staging slots, another fetches them into a request of its own, and a test lease holder drives the lender as the runtime will (lend, wait for the view, submit it as it stands, wait for quiescence, mark what arrived, release). - Pages come back byte for byte and resume_bounds covers the prompt's three whole blocks, over the in-memory store and, with TRTLLM_MOONCAKE_REAL=1, a real Mooncake pool. Each driver moves every row separately, several at once: the in-memory driver takes one row per operation, each on a thread of its own; Mooncake cuts each operation into one-row parts on two workers. - A fetch whose later blocks fault marks only the delivered prefix, so resume_bounds stops at the first gap. Its in-memory driver ignores cancels: the rows run in any order, and the cancel a later row's fault forwards could otherwise end the first row before it started. The test builds each backend itself, as a factory will: it opens the driver, takes the lender's hold, registers the regions with RegionBinding, starts BlobBackend, and gives the hold back once the backend shut down. The lender's MemoryRegion names no memory kind, which SlotRegion requires, so the test registers host-memory SlotRegions made from the lender's regions, as a runtime must until the lender's regions name their memory kind. An exception escaping a thread fails the test. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../sharing/test_blob_backend_round_trip.py | 256 ++++++++++++++++++ 1 file changed, 256 insertions(+) create mode 100644 tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py new file mode 100644 index 000000000000..1c0c9ad1d224 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""The blob backend between two real staging lenders. + +One manager publishes a request's KV blocks through its lender's staging slots; another fetches +them into a request of its own. A test lease holder drives the lender as the runtime does: lend, +wait for the view, submit it to the backend as it stands, wait until the attempt quiesces, mark +what arrived, release. The store is the in-memory one, or with ``TRTLLM_MOONCAKE_REAL=1`` a real +Mooncake pool (``TRTLLM_MOONCAKE_PROTOCOL``, default rdma). Allocates device pools. +""" + +from __future__ import annotations + +import os +import threading +from dataclasses import dataclass, replace + +import numpy as np +import pytest +import torch + +from tensorrt_llm._torch.disaggregation.base.backend import MemoryKind +from tensorrt_llm._torch.disaggregation.blob.backend import BlobBackend, BlobConfig, RegionBinding +from tensorrt_llm._torch.disaggregation.blob.keys import encode_key +from tensorrt_llm._torch.disaggregation.blob.mooncake import MooncakeDriver +from tensorrt_llm._torch.disaggregation.blob.mooncake_setup import MooncakeConfig, open_sdk +from tensorrt_llm._torch.disaggregation.testing.memory_store import ( + DEFAULT_CAPABILITIES, + Commit, + MemoryDriver, + MemoryStore, +) +from tensorrt_llm._torch.disaggregation.testing.object_trailer import Checksum, encode_trailer +from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import ( + LeaseView, + ResumeBounds, + StagingOptions, + attach_lender, +) + +# An exception that escapes a thread fails the test. +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="allocates KV cache pools"), + pytest.mark.filterwarnings("error::pytest.PytestUnhandledThreadExceptionWarning"), +] + +TPB = 32 +PROMPT = list(range(1000, 1097)) # three whole blocks and one token +END = 96 +BLOCKS = END // TPB +SOURCE, TARGET = 1, 2 +SCOPE = b"blob-backend-round-trip" +KEY_PREFIX = "round-trip" +WAIT_S = 60.0 +PROTOCOL = os.environ.get("TRTLLM_MOONCAKE_PROTOCOL", "rdma") +DEVICES = os.environ.get("TRTLLM_MOONCAKE_DEVICES", "") +ONE_ROW_PER_CALL = replace(DEFAULT_CAPABILITIES, max_rows_per_op=1) + + +class Wake: + """The ``WakeFn`` a test hands the backend, with bounded waits on what it signals.""" + + def __init__(self) -> None: + self._woken = threading.Condition() + + def __call__(self) -> None: + with self._woken: + self._woken.notify_all() + + def settled(self, attempt): + """The attempt's outcome, once it also quiesced.""" + with self._woken: + done = self._woken.wait_for( + lambda: attempt.poll() is not None and attempt.quiesced(), WAIT_S + ) + assert done, "the attempt did not settle" + return attempt.poll() + + +class Stores: + """Drivers over one store: the in-memory one, or with ``master`` a real Mooncake pool. + ``close`` closes the Mooncake clients only once every backend built over them shut down + cleanly. ``cancel_effective=False`` makes the in-memory drivers ignore a cancel.""" + + def __init__(self, master: str | None = None, cancel_effective: bool = True) -> None: + self._store = MemoryStore() if master is None else None + self._master = master + self._cancel_effective = cancel_effective + self._clients = [] + self.backends = [] + + def open_driver(self): + """A driver that moves each row in a call of its own, several at once.""" + if self._store is not None: + driver = MemoryDriver(self._store, capabilities=ONE_ROW_PER_CALL, background=True) + driver.set_cancel_effective(self._cancel_effective) + return driver + sdk = open_sdk(MooncakeConfig(self._master, 0, PROTOCOL, DEVICES, local_buffer_size=0)) + self._clients.append(sdk) + return MooncakeDriver(sdk, workers=2, rows_per_call=1) + + def commit(self, object_key: np.ndarray, data: bytes) -> None: + """Commits ``data`` straight into the in-memory store, with the driver's trailer.""" + key = encode_key(KEY_PREFIX, object_key.tobytes()) + trailer = encode_trailer(key.encode(), len(data), checksum=Checksum.CRC32, payload=data) + assert self._store.commit(key, data, trailer) is Commit.STORED + + def close(self) -> None: + if all(backend.shutdown(WAIT_S) == frozenset() for backend in self.backends): + for client in self._clients: + client.close() + + +@pytest.fixture(params=["memory", "mooncake"]) +def stores(request, tmp_path): + if request.param == "memory": + yield Stores() + return + if os.environ.get("TRTLLM_MOONCAKE_REAL") != "1": + pytest.skip("needs a real Mooncake pool; set TRTLLM_MOONCAKE_REAL=1") + from test_common.mooncake_utils import running_master_on_free_ports, running_segment_provider + + with running_master_on_free_ports(str(tmp_path), protocol=PROTOCOL) as master: + with running_segment_provider( + master.address, + str(tmp_path), + segment_size=256 << 20, + protocol=PROTOCOL, + device_name=DEVICES, + ): + opened = Stores(master=master.address) + yield opened + opened.close() + + +def attach(mgr): + return attach_lender(mgr, scope=SCOPE, staging=StagingOptions(END)) + + +def ready_view(lease, mgr, tries: int = 3): + for _ in range(tries): + mgr._stream.synchronize() + view = lease.poll() + if view is not None: + return view + raise AssertionError(f"the lease never became ready: {lease.failure}") + + +@dataclass(frozen=True) +class HostRegion: + """A lender region as the contract's ``SlotRegion``. The lender's own regions do not name their + memory kind, which ``SlotRegion`` requires; the lender stages in host memory.""" + + layout_label: str + address: int + nbytes: int + slot_bytes: int + slot_count: int + memory_kind: MemoryKind = MemoryKind.HOST + + +def build(stores: Stores, lender, wake: Wake): + """A backend over ``lender``'s regions, built on this thread: a driver, the lender's hold, the + regions registered under it, and the backend. Returns the binding, whose ``release_hold`` + gives the hold back once the backend shut down, and the backend.""" + regions = [ + HostRegion(r.layout_label, r.address, r.nbytes, r.slot_bytes, r.slot_count) + for r in lender.regions + ] + binding = RegionBinding(regions, stores.open_driver(), lender.hold_regions()) + binding.register() + backend = BlobBackend(binding, wake, BlobConfig(key_prefix=KEY_PREFIX)) + stores.backends.append(backend) + return binding, backend + + +def publish(stores, kit, mgr, select=None): + """Publishes the source request's blocks (those ``select`` keeps); returns the view lent.""" + source = kit.published(mgr, SOURCE, PROMPT) + lender, wake = attach(mgr), Wake() + binding, backend = build(stores, lender, wake) + lease = lender.lend_read(source, 0, END) + view = ready_view(lease, mgr) + submitted = view if select is None else LeaseView(tuple(g.select(select) for g in view.groups)) + outcome = wake.settled(backend.publish(submitted, None)) + assert outcome.fault is None + assert all(mask.all() for mask in outcome.delivered_masks) + lease.release() + assert backend.shutdown(WAIT_S) == frozenset() + binding.release_hold() + return source, view + + +def fetch(stores, kit, mgr): + """Fetches every block of the target request; returns the target, lender, lease, outcome, + backend and binding, the lease not yet marked.""" + target = kit.admitted(mgr, TARGET, PROMPT) + lender, wake = attach(mgr), Wake() + binding, backend = build(stores, lender, wake) + lease = lender.lend_write(target, 0, END) + view = lease.poll() + kit.fill_sentinel(mgr, target) + outcome = wake.settled(backend.fetch(view, None)) + return target, lender, lease, outcome, backend, binding + + +def test_publish_then_fetch_restores_pages(kit, real_manager, stores) -> None: + with real_manager() as mgr_a, real_manager() as mgr_b: + source, _ = publish(stores, kit, mgr_a) + target, lender, lease, outcome, backend, binding = fetch(stores, kit, mgr_b) + assert outcome.fault is None + lease.mark_arrived([np.array(mask) for mask in outcome.delivered_masks]) + mgr_b._stream.synchronize() + assert lender.resume_bounds(target) == ResumeBounds(END, 0) + for ordinal in range(BLOCKS): + got = kit.digest(kit.page(mgr_b, target, 0, ordinal)) + assert got == kit.digest(kit.page(mgr_a, source, 0, ordinal)), "a fetched page differs" + lease.release() + assert backend.shutdown(WAIT_S) == frozenset() + binding.release_hold() + + +def test_a_partial_fetch_resumes_only_up_to_its_first_undelivered_block(kit, real_manager) -> None: + # Each row's call runs on a thread of its own, in any order. A later row's fault stops the + # fetch, and the cancel it forwards could end the first row before its thread took it; the + # driver ignores it, so the first row always lands. + stores = Stores(cancel_effective=False) + with real_manager() as mgr_a, real_manager() as mgr_b: + first_block_only = np.arange(BLOCKS) == 0 + source, view = publish(stores, kit, mgr_a, select=first_block_only) + for object_key in view.groups[0].object_keys[1:]: # the later blocks are the wrong size + stores.commit(object_key, b"short") + target, lender, lease, outcome, backend, binding = fetch(stores, kit, mgr_b) + assert outcome.delivered_masks[0].tolist() == [True, False, False] + assert "object is 5 bytes" in outcome.fault + lease.mark_arrived([np.array(mask) for mask in outcome.delivered_masks]) + mgr_b._stream.synchronize() + assert lender.resume_bounds(target) == ResumeBounds(TPB, 0) # usable up to the first gap + sentinel = kit.digest(bytes([kit.SENTINEL]) * kit.DevicePages(mgr_b).page_bytes(0)) + pages = [kit.digest(kit.page(mgr_b, target, 0, ordinal)) for ordinal in range(BLOCKS)] + assert pages == [kit.digest(kit.page(mgr_a, source, 0, 0)), sentinel, sentinel] + lease.release() + assert backend.shutdown(WAIT_S) == frozenset() + binding.release_hold() From 840590063e731622bd8225e545d7fd32a08a52c5 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Tue, 27 Oct 2026 08:23:24 +0000 Subject: [PATCH 05/10] [None][fix] Revoke the Mooncake sessions left open in close, so drain makes no SDK call drain revoked the put sessions left open once it had drained, in one batch_put_session_revoke outside its budget. A slow or hung revoke then held up drain(0.0) on the backend's progress thread, and with it the revocation of the registrations and the lender's hold, although no access remained. close, the owner's last step under the blob backend's rule, now tries the revocation once before it closes the SDK client; close already waits on the master for the unmount. These sessions hold no lender memory, so nothing waits for them, and one whose revocation fails is freed by the master after put_start_release_timeout_sec. drain makes no SDK call and returns within its budget. The ledger's FOR_DRAIN becomes FOR_CLOSE, and the stats key put_sessions_left_for_close. close's docstring and open_mooncake_driver's say who closes: the caller that opened the driver, last, off the executor loop's thread. The setup and round-trip tests close their drivers that way, and the round trip keeps drivers rather than SDK clients for its fallback close. A new test holds a revoke inside the SDK and checks that drain returns at once without calling it; the session tests check that close revokes what drain left. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../_torch/disaggregation/blob/mooncake.py | 49 ++++++------ .../disaggregation/blob/mooncake_setup.py | 3 + .../blob/test_mooncake_driver.py | 79 ++++++++++++------- .../disaggregation/blob/test_mooncake_real.py | 2 + .../blob/test_mooncake_setup.py | 3 + .../sharing/test_blob_backend_round_trip.py | 27 ++++--- 6 files changed, 101 insertions(+), 62 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/blob/mooncake.py b/tensorrt_llm/_torch/disaggregation/blob/mooncake.py index 94351d1a222f..4fe7b3fd7d92 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/mooncake.py +++ b/tensorrt_llm/_torch/disaggregation/blob/mooncake.py @@ -166,7 +166,7 @@ class _Code(IntEnum): } # Session-end answers after which the end may still reach the master, even if the retry is refused. -# Such a session is left for drain and its key is never put again: the master matches an end by key +# Such a session is left for close and its key is never put again: the master matches an end by key # and client only, so a late end would mark complete this client's next object of the key. _END_OUTCOME_UNKNOWN = (_Code.RPC_FAIL, _Code.RPC_TIMEOUT, _Code.METADATA_CALL_FAILED) # Session-end answers retried once, as the master ends a session idempotently; an end call that @@ -177,7 +177,7 @@ class _Code(IntEnum): class _LeftOpen(enum.Enum): """Who ends a put session this driver left open.""" - FOR_DRAIN = "for_drain" # drain revokes it once no part runs + FOR_CLOSE = "for_close" # close revokes it, as the owner's last step TO_MASTER = "to_master" # the target may still write into its object: the master frees it @@ -303,12 +303,13 @@ class MooncakeDriver: - to the master, never revoked here, since the target may still write into the object: a scatter that raised or miscounted, a row answered TRANSFER_FAIL. The master frees them after ``put_start_release_timeout_sec``. - - for ``drain`` to revoke once no part runs: any other, such as a start that raised or + - for ``close`` to revoke, once no put runs: any other, such as a start that raised or miscounted, an end whose outcome is unknown (``_END_OUTCOME_UNKNOWN``), a revoke that failed, or a put stopped by a raise. Args: - store: An SDK store whose ``setup`` returned 0. + store: An SDK store whose ``setup`` returned 0. The driver takes it over: ``close`` + closes it. workers: The most threads running SDK calls; each blocks while its call waits on the store. rows_per_call: The most rows of one part, so of each SDK call. @@ -471,12 +472,9 @@ def probe( def drain(self, budget_s: float) -> bool: """Refuses new work for good, cancels the rows no worker took, ends the worker threads, and waits up to ``budget_s`` for the running parts to end and every worker thread to exit. - Returns False at once while an operation is unprovable. - - Each drain that returns True also tries once to revoke the sessions left for drain: no put - runs any more, so none is in use. They hold no lender memory, so they never change the - answer; one whose revocation fails stays for the next drain, or until the master discards - it. Sessions left to the master are never revoked here.""" + Returns False at once while an operation is unprovable. Makes no SDK call, so it returns + within ``budget_s``: the sessions left for close hold no lender memory, and ``close`` + revokes them.""" with self._lock: self._require_open() self._draining = True @@ -495,19 +493,14 @@ def drain(self, budget_s: float) -> bool: ) if drained: self._drained = True - left = self._sessions_left_open.items() - to_revoke = [key for key, by in left if by is _LeftOpen.FOR_DRAIN] if drained else [] - if to_revoke: - codes = self._metadata_call(self._store.batch_put_session_revoke, to_revoke) - with self._lock: - for key, code in zip(to_revoke, codes): - if _session_gone(code): - del self._sessions_left_open[key] - return drained + return drained def close(self) -> None: - """Closes the SDK client once a drain returned True and every region is unregistered. - Never called in production: it holds the GIL through unmount RPCs; process exit frees it.""" + """The owner's last step, once a drain returned True and every region is unregistered: + tries once to revoke the sessions left for close, then closes the SDK client. Both wait on + the master and the unmount with the GIL held, so the owner calls it off the executor loop's + thread. A session whose revocation fails is freed by the master after + ``put_start_release_timeout_sec``; sessions left to the master are never revoked here.""" with self._lock: if self._closed: return @@ -516,6 +509,14 @@ def close(self) -> None: if any(r.live for r in self._registrations): raise DriverRejected("close needs every registration unregistered") self._closed = True + left = self._sessions_left_open.items() + to_revoke = [key for key, by in left if by is _LeftOpen.FOR_CLOSE] + if to_revoke: # no put runs any more, so none of them is in use + codes = self._metadata_call(self._store.batch_put_session_revoke, to_revoke) + with self._lock: + for key, code in zip(to_revoke, codes): + if _session_gone(code): + del self._sessions_left_open[key] self._store.close() def stats(self) -> dict[str, int]: @@ -528,7 +529,7 @@ def stats(self) -> dict[str, int]: for reason, r in _SKIP_RESULTS.items() }, "ops_unprovable": self._unprovable_op_count, - "put_sessions_left_for_drain": left.count(_LeftOpen.FOR_DRAIN), + "put_sessions_left_for_close": left.count(_LeftOpen.FOR_CLOSE), "put_sessions_left_to_master": left.count(_LeftOpen.TO_MASTER), } @@ -782,11 +783,11 @@ def _release_keys_and_record_left_open( self, reserved: list[str], closed: set[str], to_master: set[str] ) -> None: """Ends the reservation of ``reserved`` and records each session not ``closed`` as left - open: for the master when its key is in ``to_master``, otherwise for ``drain``.""" + open: for the master when its key is in ``to_master``, otherwise for ``close``.""" left: dict[str, _LeftOpen] = {} if not closed.issuperset(reserved): # usually every session ended left = { - key: _LeftOpen.TO_MASTER if key in to_master else _LeftOpen.FOR_DRAIN + key: _LeftOpen.TO_MASTER if key in to_master else _LeftOpen.FOR_CLOSE for key in reserved if key not in closed } diff --git a/tensorrt_llm/_torch/disaggregation/blob/mooncake_setup.py b/tensorrt_llm/_torch/disaggregation/blob/mooncake_setup.py index b09acd306ecc..5683fa30cd8d 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/mooncake_setup.py +++ b/tensorrt_llm/_torch/disaggregation/blob/mooncake_setup.py @@ -219,6 +219,9 @@ def open_mooncake_driver( """Opens the Mooncake driver a store config describes, once both its sections are read: opening waits for the master and mounts this process's segment, so a bad config fails before that. + The caller owns the driver and closes it last, with ``MooncakeDriver.close``, which also closes + the SDK client. + Args: config: A store config, as ``load_store_config`` takes it. read_blob_section: Reads the ``blob`` section (for the blob backend, into a diff --git a/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py index 34bd94a5fe78..0788d1247712 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py +++ b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py @@ -756,7 +756,7 @@ def test_put_session_start_answers_map_to_row_statuses(case: str) -> None: PUT_WRITE_ANSWERS = { # case: (the scatter's answer, the end's answer (None: real), status, fault, the session's end: - # ended, revoked, left_for_drain or left_to_master) + # ended, revoked, left_for_close or left_to_master) "written": (None, None, "ready", None, "ended"), "short_write": (5, None, "fault", "put ranges: moved 5 of 8 bytes", "revoked"), # A TCP target may still write into the object's space. @@ -765,7 +765,7 @@ def test_put_session_start_answers_map_to_row_statuses(case: str) -> None: "end_refused": (None, -704, "fault", "put session end: OBJECT_NOT_FOUND(-704)", "revoked"), "end_not_owner": (None, -601, "fault", "put session end: ILLEGAL_CLIENT(-601)", "revoked"), # The end may still reach the master. - "end_timed_out": (None, -901, "fault", "put session end: RPC_TIMEOUT(-901)", "left_for_drain"), + "end_timed_out": (None, -901, "fault", "put session end: RPC_TIMEOUT(-901)", "left_for_close"), } @@ -783,7 +783,7 @@ def test_put_ranges_and_end_answers_map_to_row_statuses(case: str) -> None: assert (bench.fake.pool.objects.get(key("A")) == b"aaaaaaaa") == (status == "ready") assert bench.keys_of(PUT_REVOKE) == ([key("A")] if session == "revoked" else []) stats = bench.driver.stats() - assert stats["put_sessions_left_for_drain"] == (session == "left_for_drain") + assert stats["put_sessions_left_for_close"] == (session == "left_for_close") assert stats["put_sessions_left_to_master"] == (session == "left_to_master") @@ -824,7 +824,7 @@ class SessionScenario(NamedTuple): rows: str # each row ready (R) or faulted (F) access_ends: bool = True revoked: str = "" # rows whose session the put revokes, or tries to - left_for_drain: str = "" # rows whose session is left open for drain to revoke + left_for_close: str = "" # rows whose session is left open for close to revoke left_to_master: str = "" # rows whose session is left open for the master alone to free @@ -836,32 +836,32 @@ class SessionScenario(NamedTuple): "end_retried": SessionScenario({PUT_END: [[-900, None, None], [None]]}, "RRR"), # Ends that may still reach the master. "end_fails_twice": SessionScenario( - {PUT_END: [[-901, None, None], [-900]]}, "FRR", left_for_drain="A" + {PUT_END: [[-901, None, None], [-900]]}, "FRR", left_for_close="A" ), "end_retry_refused": SessionScenario( - {PUT_END: [[-901, None, None], [-704]]}, "FRR", left_for_drain="A" + {PUT_END: [[-901, None, None], [-704]]}, "FRR", left_for_close="A" ), "end_retry_raises": SessionScenario( - {PUT_END: [[-901, None, None], RuntimeError("x")]}, "FRR", left_for_drain="A" + {PUT_END: [[-901, None, None], RuntimeError("x")]}, "FRR", left_for_close="A" ), "end_raises": SessionScenario( - {PUT_END: [RuntimeError("rpc lost")]}, "FFF", left_for_drain="ABC" + {PUT_END: [RuntimeError("rpc lost")]}, "FFF", left_for_close="ABC" ), "end_committed_then_raised": SessionScenario( - {PUT_END: [reply_lost]}, "FFF", left_for_drain="ABC" + {PUT_END: [reply_lost]}, "FFF", left_for_close="ABC" ), "end_and_revoke_fail": SessionScenario( {PUT_END: [[-704, None, None]], PUT_REVOKE: [[-900]]}, "FRR", revoked="A", - left_for_drain="A", + left_for_close="A", ), # The SDK may have opened any of the sessions. "start_raises": SessionScenario( - {PUT_START: [RuntimeError("rpc lost")]}, "FFF", left_for_drain="ABC" + {PUT_START: [RuntimeError("rpc lost")]}, "FFF", left_for_close="ABC" ), "start_miscounts": SessionScenario( - {PUT_START: [lambda answers: answers[:-1]]}, "FFF", left_for_drain="ABC" + {PUT_START: [lambda answers: answers[:-1]]}, "FFF", left_for_close="ABC" ), # The scatter may still be running: nothing is ended or revoked. "ranges_raise": SessionScenario( @@ -889,35 +889,39 @@ def test_every_started_put_session_is_ended_revoked_or_left_open(name: str) -> N assert bench.keys_of(PUT_REVOKE) == [key(n) for n in scenario.revoked] settled = Counter(bench.fake.keys_answered(PUT_END, 0) + bench.keys_of(PUT_REVOKE)) assert set(settled.values()) <= {1}, f"a session ended or revoked twice: {settled}" - for_drain = {key(n) for n in scenario.left_for_drain} + for_close = {key(n) for n in scenario.left_for_close} to_master = {key(n) for n in scenario.left_to_master} stats = bench.driver.stats() - assert stats["put_sessions_left_for_drain"] == len(for_drain) + assert stats["put_sessions_left_for_close"] == len(for_close) assert stats["put_sessions_left_to_master"] == len(to_master) # Every session the SDK still holds is left open; the SDK may have dropped one the driver # could not see end, as when an end landed but its reply did not. - assert bench.fake.put_sessions <= for_drain | to_master + assert bench.fake.put_sessions <= for_close | to_master # Only a row written whole may be ended: its object becomes readable as it stands. written = bench.fake.keys_answered(PUT_RANGES, SLOT) assert set(bench.keys_of(PUT_END)) <= set(written) ready = [key(n) for n, r in zip("ABC", results) if r.status is RowStatus.READY] assert set(ready) == set(bench.fake.keys_answered(PUT_END, 0)) - # A drain revokes the sessions left for it, never those left to the master. + # Close revokes the sessions left for it, never those left to the master; drain revokes none. if scenario.access_ends: + revocations = len(bench.keys_of(PUT_REVOKE)) assert bench.driver.drain(WAIT_S) is True - assert bench.driver.stats()["put_sessions_left_for_drain"] == 0 + assert len(bench.keys_of(PUT_REVOKE)) == revocations + bench.driver.unregister(bench.handle) + bench.driver.close() + assert bench.driver.stats()["put_sessions_left_for_close"] == 0 assert bench.driver.stats()["put_sessions_left_to_master"] == len(to_master) assert bench.fake.put_sessions == to_master assert not to_master & set(bench.keys_of(PUT_REVOKE)) -def test_sessions_left_open_are_revoked_only_in_drain() -> None: +def test_sessions_left_open_are_revoked_only_in_close() -> None: bench = Bench() bench.write(0, b"aaaaaaaa") bench.fake.script(PUT_END, [-704]) bench.fake.script(PUT_REVOKE, [-900]) bench.run("put", "A") - assert bench.driver.stats()["put_sessions_left_for_drain"] == 1 + assert bench.driver.stats()["put_sessions_left_for_close"] == 1 bench.write(1, b"bbbbbbbb") bench.store(C=b"cccccccc") bench.run("put", "X", "B") @@ -927,16 +931,37 @@ def test_sessions_left_open_are_revoked_only_in_drain() -> None: assert summary(results) == [("not_ready", None)] # and no complete object exists assert bench.keys_of(PUT_START).count(key("A")) == 1 - bench.fake.script(PUT_REVOKE, [-900]) assert bench.driver.drain(WAIT_S) is True # sessions left open hold no lender memory + assert bench.keys_of(PUT_REVOKE) == [key("A")] + bench.driver.unregister(bench.handle) + bench.fake.script(PUT_REVOKE, [-900]) + bench.driver.close() assert bench.keys_of(PUT_REVOKE) == [key("A")] * 2 - assert bench.driver.stats()["put_sessions_left_for_drain"] == 1 # retried by the next drain - assert bench.driver.drain(0.0) is True - assert bench.driver.stats()["put_sessions_left_for_drain"] == 0 + assert bench.driver.stats()["put_sessions_left_for_close"] == 1 # the master frees it later + assert len(bench.fake.calls_to("close")) == 1 + + +def test_drain_returns_at_once_while_a_revocation_would_block() -> None: + bench = Bench() + bench.fill_sources("put", "A") + bench.fake.script(PUT_END, [-901]) + bench.fake.script(PUT_END, [-901]) # the end and its retry go unanswered + bench.run("put", "A") + assert bench.driver.stats()["put_sessions_left_for_close"] == 1 + gate = bench.fake.block(PUT_REVOKE) + started = time.monotonic() + assert bench.driver.drain(WAIT_S) is True # once the idle workers exit + assert time.monotonic() - started < PROMPT_S + assert not gate.entered.is_set() + gate.release.set() + bench.driver.unregister(bench.handle) + bench.driver.close() + assert bench.keys_of(PUT_REVOKE) == [key("A")] + assert bench.driver.stats()["put_sessions_left_for_close"] == 0 assert bench.fake.put_sessions == set() -def test_a_put_whose_worker_thread_exits_leaves_its_session_open_for_drain( +def test_a_put_whose_worker_thread_exits_leaves_its_session_open_for_close( thread_escapes: list[BaseException], ) -> None: bench = Bench() @@ -946,7 +971,7 @@ def test_a_put_whose_worker_thread_exits_leaves_its_session_open_for_drain( assert summary(results) == [("fault", "the worker stopped before the row was decided")] assert not access_ended assert key("A") in bench.fake.put_sessions - assert bench.driver.stats()["put_sessions_left_for_drain"] == 1 + assert bench.driver.stats()["put_sessions_left_for_close"] == 1 results, _ = bench.run("put", "A") # the driver holds A's session: it starts no other assert summary(results) == [("not_ready", None)] assert bench.keys_of(PUT_START).count(key("A")) == 1 @@ -975,7 +1000,7 @@ def test_a_key_another_put_of_this_driver_holds_is_not_ready_without_asking_the_ assert summary(results) == [("ready", None)] # it exists already, committed -@pytest.mark.parametrize("left", ["for_drain", "to_master"]) +@pytest.mark.parametrize("left", ["for_close", "to_master"]) def test_a_key_whose_session_was_left_open_is_ready_once_its_object_exists(left: str) -> None: bench = Bench() bench.fill_sources("put", "A") @@ -1123,7 +1148,7 @@ def broken(*_): assert summary(results) == [("fault", f"{direction} failed: ValueError('mapping bug')")] assert not access_ended # A put's session may still be open: the driver keeps holding the key. - assert bench.driver.stats()["put_sessions_left_for_drain"] == (direction == "put") + assert bench.driver.stats()["put_sessions_left_for_close"] == (direction == "put") # Admission, registration, concurrency. diff --git a/tests/unittest/_torch/disaggregation/blob/test_mooncake_real.py b/tests/unittest/_torch/disaggregation/blob/test_mooncake_real.py index 840480e8bae9..01707545b1b2 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_mooncake_real.py +++ b/tests/unittest/_torch/disaggregation/blob/test_mooncake_real.py @@ -154,6 +154,8 @@ class Clients: ``close`` closes the clients only once every endpoint's shutdown returns no labels: a backend thread still inside the SDK would crash on a closed client, so a failed test leaves them open. + A closed endpoint's driver has closed its client already, and closing a client twice does + nothing. """ def __init__(self, master: str) -> None: diff --git a/tests/unittest/_torch/disaggregation/blob/test_mooncake_setup.py b/tests/unittest/_torch/disaggregation/blob/test_mooncake_setup.py index 82f8b508680e..ca4e47eede9a 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_mooncake_setup.py +++ b/tests/unittest/_torch/disaggregation/blob/test_mooncake_setup.py @@ -226,6 +226,8 @@ def driver_recording_its_tuning(store, **tuning) -> MooncakeDriver: assert store.calls_to("register_buffer") driver.unregister(handle) assert driver.drain(DRAIN_S) # its worker threads exit + driver.close() # the caller owns the driver, and closing it closes the client + assert len(store.calls_to("close")) == 1 open_sdk(MooncakeConfig(MASTER, 1 << 30, "tcp", "eth9", "10.9.9.9", "etcd://m:2379", 0)) assert sdk.stores[1].calls_to("setup")[0].kwargs == { @@ -337,6 +339,7 @@ def test_only_pinned_sdk_builds_are_opened( driver, _ = open_mooncake_driver(store_config(), read_blob) assert len(sdk.stores) == 1 assert driver.drain(DRAIN_S) # its worker threads exit + driver.close() return with pytest.raises(ImportError, match=f"comes from .*{refused}.*follows 0.3.13 of"): open_mooncake_driver(store_config(), read_blob) diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py index 1c0c9ad1d224..3df4ba91b8ce 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py @@ -90,15 +90,16 @@ def settled(self, attempt): class Stores: - """Drivers over one store: the in-memory one, or with ``master`` a real Mooncake pool. - ``close`` closes the Mooncake clients only once every backend built over them shut down - cleanly. ``cancel_effective=False`` makes the in-memory drivers ignore a cancel.""" + """Drivers over one store: the in-memory one, or with ``master`` a real Mooncake pool. A test + closes each driver it opened as its owner does, last; ``close`` closes those a failed test left + open, only once every backend built over them shut down cleanly. ``cancel_effective=False`` + makes the in-memory drivers ignore a cancel.""" def __init__(self, master: str | None = None, cancel_effective: bool = True) -> None: self._store = MemoryStore() if master is None else None self._master = master self._cancel_effective = cancel_effective - self._clients = [] + self._drivers = [] self.backends = [] def open_driver(self): @@ -106,10 +107,11 @@ def open_driver(self): if self._store is not None: driver = MemoryDriver(self._store, capabilities=ONE_ROW_PER_CALL, background=True) driver.set_cancel_effective(self._cancel_effective) - return driver - sdk = open_sdk(MooncakeConfig(self._master, 0, PROTOCOL, DEVICES, local_buffer_size=0)) - self._clients.append(sdk) - return MooncakeDriver(sdk, workers=2, rows_per_call=1) + else: + sdk = open_sdk(MooncakeConfig(self._master, 0, PROTOCOL, DEVICES, local_buffer_size=0)) + driver = MooncakeDriver(sdk, workers=2, rows_per_call=1) + self._drivers.append(driver) + return driver def commit(self, object_key: np.ndarray, data: bytes) -> None: """Commits ``data`` straight into the in-memory store, with the driver's trailer.""" @@ -119,8 +121,8 @@ def commit(self, object_key: np.ndarray, data: bytes) -> None: def close(self) -> None: if all(backend.shutdown(WAIT_S) == frozenset() for backend in self.backends): - for client in self._clients: - client.close() + for driver in self._drivers: + driver.close() @pytest.fixture(params=["memory", "mooncake"]) @@ -174,7 +176,7 @@ class HostRegion: def build(stores: Stores, lender, wake: Wake): """A backend over ``lender``'s regions, built on this thread: a driver, the lender's hold, the regions registered under it, and the backend. Returns the binding, whose ``release_hold`` - gives the hold back once the backend shut down, and the backend.""" + gives the hold back once the backend shut down, before its driver is closed, and the backend.""" regions = [ HostRegion(r.layout_label, r.address, r.nbytes, r.slot_bytes, r.slot_count) for r in lender.regions @@ -200,6 +202,7 @@ def publish(stores, kit, mgr, select=None): lease.release() assert backend.shutdown(WAIT_S) == frozenset() binding.release_hold() + binding.driver.close() return source, view @@ -230,6 +233,7 @@ def test_publish_then_fetch_restores_pages(kit, real_manager, stores) -> None: lease.release() assert backend.shutdown(WAIT_S) == frozenset() binding.release_hold() + binding.driver.close() def test_a_partial_fetch_resumes_only_up_to_its_first_undelivered_block(kit, real_manager) -> None: @@ -254,3 +258,4 @@ def test_a_partial_fetch_resumes_only_up_to_its_first_undelivered_block(kit, rea lease.release() assert backend.shutdown(WAIT_S) == frozenset() binding.release_hold() + binding.driver.close() From 321fc530e8d54ad7df604d8c51db088fbe798812 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Tue, 27 Oct 2026 08:23:30 +0000 Subject: [PATCH 06/10] [None][test] Pin a repeated object key over Mooncake and say where NOT_READY departs The blob backend cuts a group at a repeated object key, so no operation names a key twice; the driver keys a part's results by store key and relies on that. A test over the fake SDK publishes and fetches a view naming one key twice: two calls, no fault, both slots fetched. 0.3.13 cannot always confirm the hold that NOT_READY stands for: -703 also answers a complete replica it cannot read, batch_is_exist answers 0 for one too, and a session left open after a start that raised may not exist. The driver keeps NOT_READY for them, since the usual cause is a write in progress and a FAULT would stop every fetch that races a writer and count toward isolation; _SKIP_RESULTS now says so. Comments also say why a failed register_buffer certifies nothing was registered, and when the last worker's on_change may come after a drain returned True. The round trip's open_driver says how each driver moves its rows one by one. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../_torch/disaggregation/blob/mooncake.py | 13 ++++++- .../blob/test_mooncake_driver.py | 37 +++++++++++++++++++ .../sharing/test_blob_backend_round_trip.py | 3 +- 3 files changed, 51 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/blob/mooncake.py b/tensorrt_llm/_torch/disaggregation/blob/mooncake.py index 4fe7b3fd7d92..e78cb15601a3 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/mooncake.py +++ b/tensorrt_llm/_torch/disaggregation/blob/mooncake.py @@ -154,6 +154,12 @@ class _Code(IntEnum): # Why a row was skipped (nothing moved and nothing went wrong), and the result it gets; ``stats()`` # counts skipped rows by both, and a part's warning names its first row that faulted or was # skipped, with the skip's reason. +# +# NOT_READY here departs from the protocol's "only for a hold the store confirmed": 0.3.13 cannot +# tell a write in progress from a complete replica it cannot read. -703 answers both, and so does +# batch_is_exist's 0 after a -705; a session left open after a start that raised may not exist. +# The usual cause is a write in progress, which as a FAULT would stop every fetch that races a +# writer and count toward isolation; a replica that stays unreadable costs only its rows. _SKIP_RESULTS = { # get -703: no complete replica, as while the object is written or once its replica's # endpoint went stale @@ -407,7 +413,9 @@ def register(self, address: int, nbytes: int, layout: RegionLayout) -> _Registra raise RegistrationError(f"register_buffer raised {error!r}", registration) from error # The SDK undoes a partial registration before it answers (RdmaTransport:: # registerLocalMemoryInternal), so a non-OK answer leaves nothing registered; a rollback - # that itself fails is only logged by the SDK and cannot be seen here. + # that itself fails is only logged by the SDK and cannot be seen here. Handing back the + # handle instead would keep the lender's hold for good: unregistering what the SDK never + # kept fails. if code != _Code.OK: raise RegistrationError(f"register_buffer returned {_describe(code)}") self._add(registration) @@ -606,6 +614,9 @@ def _worker_loop(self) -> None: if not self._workers_alive: done = self._end_untaken_rows(self._queued_ops(), _NO_WORKER_LEFT) self._lock.notify_all() # drain waits for every worker thread to exit + # Called after this worker stopped counting, so a drain may return True first. Rows + # are left queued only once every worker died of a BaseException that left no + # operation unprovable, which only an on_change raising against the protocol does. for op in done: self._call_on_change(op) diff --git a/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py index 0788d1247712..3e142723fba1 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py +++ b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py @@ -1378,6 +1378,43 @@ def test_backend_round_trip_over_the_fake_store() -> None: reader.close() +@pytest.mark.parametrize("direction", ["publish", "fetch"]) +def test_a_view_naming_one_key_twice_goes_in_two_calls_and_never_faults(direction: str) -> None: + """The backend cuts a group at a repeated object key, so no operation names a key twice. With + one worker the second put finds the object the first one committed; running alongside it, it + would find the key held by this driver's put and be NOT_READY, never a fault.""" + fake = FakeMooncakeStore() + slots = ["S0", "S1", "S2"] + endpoint = Endpoint(MooncakeDriver(fake, workers=1), slots) + keys = "ABA" + payload = {name: name.encode() * SLOT for name in set(keys)} + view = CpuView( + [Rows(0, np.stack([object_key(name) for name in keys]), endpoint.region.slots(0, 1, 2), 0)] + ) + if direction == "publish": + for slot, name in zip(slots, keys): + endpoint.write(slot, payload[name]) + attempt = endpoint.backend.publish(view, None) + else: + for name, data in payload.items(): + fake.pool.objects[store_key(name)] = data + attempt = endpoint.backend.fetch(view, None) + outcome = endpoint.wait_outcome(attempt) + assert outcome.fault is None + assert outcome.delivered_masks[0].tolist() == [True, True, True] + first_call = PUT_START if direction == "publish" else GET + assert [call.args[0] for call in fake.calls_to(first_call)] == [ + [store_key("A"), store_key("B")], + [store_key("A")], + ] + if direction == "publish": + assert {name: fake.pool.objects[store_key(name)] for name in "AB"} == payload + else: + assert [endpoint.read(slot) for slot in slots] == [payload[name] for name in keys] + endpoint.wait_quiesced(attempt) + endpoint.close() + + @pytest.fixture def work_per_row_or_thread(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, int]]: """Every key or row a ``RowBatch`` builds and every thread started, each with the thread that diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py index 3df4ba91b8ce..677250a1045f 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py @@ -103,7 +103,8 @@ def __init__(self, master: str | None = None, cancel_effective: bool = True) -> self.backends = [] def open_driver(self): - """A driver that moves each row in a call of its own, several at once.""" + """A driver that moves every row on its own, several at once: the in-memory one takes one + row per operation, Mooncake cuts each operation into one-row parts for two workers.""" if self._store is not None: driver = MemoryDriver(self._store, capabilities=ONE_ROW_PER_CALL, background=True) driver.set_cancel_effective(self._cancel_effective) From 83496c94f31e39d6d01177c99e455ecf176d42a8 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Tue, 27 Oct 2026 08:46:41 +0000 Subject: [PATCH 07/10] [None][fix] Close the Mooncake client even when its revocation raises, and pin the order What the review of the drain and close change found: - close closes the SDK client in a finally, so a BaseException from the revocation no longer leaves the client open behind a driver that already counts as closed. - Nothing pinned that close revokes before it closes the client. A closed 0.3.13 client answers every batch call with an empty list, which the driver would take as a failed revocation. The fake SDK now answers so too, and the session tests then fail if the order flips. - close's docstring: the revocation waits on the master without the GIL; only closing the client holds it through the unmount. - A put stopped by a raise is unprovable, so close never comes for its sessions; the class docstring says the master frees them. - The NOT_READY note names its one lasting case: a key whose left-open session the master freed stays NOT_READY, as it is never put again. - The worker-exit comment says when a drain returns True before the last worker's on_change, instead of when rows are left queued. - test_drain_leaves_the_session_of_an_unprovable_put_to_the_master is now test_an_unprovable_put_leaves_its_session_to_the_master, since drain revokes nothing. - The round trip's fallback close skips a driver a failed build never drained, so the clients after it still close, and its open_driver says Mooncake runs puts one at a time. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../_torch/disaggregation/blob/mooncake.py | 42 +++++++++++-------- .../disaggregation/blob/fake_mooncake.py | 4 ++ .../blob/test_mooncake_driver.py | 18 +++++++- .../sharing/test_blob_backend_round_trip.py | 10 +++-- 4 files changed, 52 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/blob/mooncake.py b/tensorrt_llm/_torch/disaggregation/blob/mooncake.py index e78cb15601a3..22634e404f94 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/mooncake.py +++ b/tensorrt_llm/_torch/disaggregation/blob/mooncake.py @@ -157,9 +157,11 @@ class _Code(IntEnum): # # NOT_READY here departs from the protocol's "only for a hold the store confirmed": 0.3.13 cannot # tell a write in progress from a complete replica it cannot read. -703 answers both, and so does -# batch_is_exist's 0 after a -705; a session left open after a start that raised may not exist. -# The usual cause is a write in progress, which as a FAULT would stop every fetch that races a -# writer and count toward isolation; a replica that stays unreadable costs only its rows. +# batch_is_exist's 0 after a -705. A session left open after a start that raised may not exist, +# and once the master frees one after put_start_release_timeout_sec, its key stays NOT_READY for +# this driver's life, as it is never put again. The usual cause is a write in progress, which as +# a FAULT would stop every fetch that races a writer and count toward isolation; a replica that +# stays unreadable costs only its rows. _SKIP_RESULTS = { # get -703: no complete replica, as while the object is written or once its replica's # endpoint went stale @@ -311,7 +313,8 @@ class MooncakeDriver: after ``put_start_release_timeout_sec``. - for ``close`` to revoke, once no put runs: any other, such as a start that raised or miscounted, an end whose outcome is unknown (``_END_OUTCOME_UNKNOWN``), a revoke that - failed, or a put stopped by a raise. + failed, or a put stopped by a raise. A put stopped by a raise is unprovable, so no drain + returns True and ``close`` never comes: the master frees its sessions. Args: store: An SDK store whose ``setup`` returned 0. The driver takes it over: ``close`` @@ -505,10 +508,11 @@ def drain(self, budget_s: float) -> bool: def close(self) -> None: """The owner's last step, once a drain returned True and every region is unregistered: - tries once to revoke the sessions left for close, then closes the SDK client. Both wait on - the master and the unmount with the GIL held, so the owner calls it off the executor loop's - thread. A session whose revocation fails is freed by the master after - ``put_start_release_timeout_sec``; sessions left to the master are never revoked here.""" + tries once to revoke the sessions left for close, then closes the SDK client, which takes + no call after that. The revocation waits on the master, and closing holds the GIL through + the unmount, so the owner calls it off the executor loop's thread. A session whose + revocation fails is freed by the master after ``put_start_release_timeout_sec``; sessions + left to the master are never revoked here.""" with self._lock: if self._closed: return @@ -519,13 +523,15 @@ def close(self) -> None: self._closed = True left = self._sessions_left_open.items() to_revoke = [key for key, by in left if by is _LeftOpen.FOR_CLOSE] - if to_revoke: # no put runs any more, so none of them is in use - codes = self._metadata_call(self._store.batch_put_session_revoke, to_revoke) - with self._lock: - for key, code in zip(to_revoke, codes): - if _session_gone(code): - del self._sessions_left_open[key] - self._store.close() + try: + if to_revoke: # no put runs any more, so none of them is in use + codes = self._metadata_call(self._store.batch_put_session_revoke, to_revoke) + with self._lock: + for key, code in zip(to_revoke, codes): + if _session_gone(code): + del self._sessions_left_open[key] + finally: + self._store.close() def stats(self) -> dict[str, int]: """A snapshot for logs and tests; not part of the driver protocol.""" @@ -614,9 +620,9 @@ def _worker_loop(self) -> None: if not self._workers_alive: done = self._end_untaken_rows(self._queued_ops(), _NO_WORKER_LEFT) self._lock.notify_all() # drain waits for every worker thread to exit - # Called after this worker stopped counting, so a drain may return True first. Rows - # are left queued only once every worker died of a BaseException that left no - # operation unprovable, which only an on_change raising against the protocol does. + # Called after this worker stopped counting. A drain returns True before them only + # once every worker died of a BaseException that left no operation unprovable, which + # only an on_change raising against the protocol does. for op in done: self._call_on_change(op) diff --git a/tests/unittest/_torch/disaggregation/blob/fake_mooncake.py b/tests/unittest/_torch/disaggregation/blob/fake_mooncake.py index 0bd9469f577a..53771b65b1d6 100644 --- a/tests/unittest/_torch/disaggregation/blob/fake_mooncake.py +++ b/tests/unittest/_torch/disaggregation/blob/fake_mooncake.py @@ -100,6 +100,7 @@ def __init__(self, pool: FakePool | None = None) -> None: self.registered: dict[int, int] = {} self.put_sessions: set[str] = set() self.setup_answer = 0 + self.closed = False self._lock = threading.Lock() self._steps: dict[str, deque] = defaultdict(deque) self._gates: dict[str, deque[Gate]] = defaultdict(deque) @@ -139,6 +140,7 @@ def setup(self, *args, **kwargs) -> int: def close(self) -> int: self._record("close", (), {}) + self.closed = True return 0 def register_buffer(self, buffer_ptr: int, size: int) -> int: @@ -281,6 +283,8 @@ def _record(self, method: str, args: tuple, kwargs: dict) -> SdkCall: def _answer(self, method, args, kwargs, count, real: Callable[[int], int]) -> list[int]: call = self._record(method, args, kwargs) + if self.closed and method.startswith("batch_"): + return [] # a closed 0.3.13 client answers every batch call with no answer at all with self._lock: gate = self._gates[method].popleft() if self._gates[method] else None step = self._steps[method].popleft() if self._steps[method] else None diff --git a/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py index 3e142723fba1..4693f19e33fc 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py +++ b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py @@ -961,6 +961,22 @@ def test_drain_returns_at_once_while_a_revocation_would_block() -> None: assert bench.fake.put_sessions == set() +def test_close_closes_the_client_even_when_the_revocation_raises() -> None: + bench = Bench() + bench.fill_sources("put", "A") + bench.fake.script(PUT_END, [-901]) + bench.fake.script(PUT_END, [-901]) # the end and its retry go unanswered + bench.run("put", "A") + assert bench.driver.drain(WAIT_S) is True + bench.driver.unregister(bench.handle) + bench.fake.script(PUT_REVOKE, SystemExit()) # escapes every handler + with pytest.raises(SystemExit): + bench.driver.close() + assert len(bench.fake.calls_to("close")) == 1 + bench.driver.close() # and is done + assert len(bench.fake.calls_to("close")) == 1 + + def test_a_put_whose_worker_thread_exits_leaves_its_session_open_for_close( thread_escapes: list[BaseException], ) -> None: @@ -1042,7 +1058,7 @@ def test_close_needs_a_drain_that_returned_true() -> None: assert len(bench.fake.calls_to("close")) == 1 -def test_drain_leaves_the_session_of_an_unprovable_put_to_the_master() -> None: +def test_an_unprovable_put_leaves_its_session_to_the_master() -> None: bench = Bench() bench.fake.script(PUT_RANGES, RuntimeError("bad pointer")) _, access_ended = bench.run("put", "A") diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py index 677250a1045f..1cff9485d5a5 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py @@ -23,6 +23,7 @@ from __future__ import annotations +import contextlib import os import threading from dataclasses import dataclass, replace @@ -36,6 +37,7 @@ from tensorrt_llm._torch.disaggregation.blob.keys import encode_key from tensorrt_llm._torch.disaggregation.blob.mooncake import MooncakeDriver from tensorrt_llm._torch.disaggregation.blob.mooncake_setup import MooncakeConfig, open_sdk +from tensorrt_llm._torch.disaggregation.blob.store import DriverRejected from tensorrt_llm._torch.disaggregation.testing.memory_store import ( DEFAULT_CAPABILITIES, Commit, @@ -103,8 +105,9 @@ def __init__(self, master: str | None = None, cancel_effective: bool = True) -> self.backends = [] def open_driver(self): - """A driver that moves every row on its own, several at once: the in-memory one takes one - row per operation, Mooncake cuts each operation into one-row parts for two workers.""" + """A driver that moves every row on its own: the in-memory one takes one row per operation, + each on a thread of its own; Mooncake cuts each operation into one-row parts for two + workers, which run gets two at once and puts one at a time.""" if self._store is not None: driver = MemoryDriver(self._store, capabilities=ONE_ROW_PER_CALL, background=True) driver.set_cancel_effective(self._cancel_effective) @@ -123,7 +126,8 @@ def commit(self, object_key: np.ndarray, data: bytes) -> None: def close(self) -> None: if all(backend.shutdown(WAIT_S) == frozenset() for backend in self.backends): for driver in self._drivers: - driver.close() + with contextlib.suppress(DriverRejected): # one a failed build never drained + driver.close() @pytest.fixture(params=["memory", "mooncake"]) From aaeeb4cdedeb3f54364cc511e8693b97200771ba Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Wed, 28 Oct 2026 10:03:09 +0000 Subject: [PATCH 08/10] [None][fix] Widen NOT_READY to what the store reports, and decline a put of a key left open What the review of the NOT_READY departure settled: - RowStatus.NOT_READY now covers an object of the key that the store reports without a readable complete replica, as while a write has not finished or when no complete replica can be read, with the row memory of a get left untouched; and another put of the same driver, still running, that writes the key. RPC, transfer and integrity errors stay FAULT. README rule 4 says the same, and both say these statuses charge the row no fault rather than that nothing went wrong. - So Mooncake's first-read -703 and a put start's -705 with no complete object are NOT_READY under the rule itself, and the departure note in mooncake.py goes. - A put of a key whose session this driver left open, with no complete object behind it, is DECLINED with its own reason and counter (rows_declined_held_by_session_left_open), instead of NOT_READY under held_by_own_put. The driver never writes such a key again, and the put that left the session open already concluded FAULT, so a FAULT here would count that one failure again on every later put of the key. Beyond a store's own refusal, RowStatus.DECLINED covers only this case. - held_by_own_put keeps its name and counter but now means only another put of this driver writing the key now. - The probe docstring follows NOT_READY: a key a get would find NOT_READY is NOT_READY there too. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../_torch/disaggregation/blob/README.md | 5 +-- .../_torch/disaggregation/blob/mooncake.py | 34 ++++++++----------- .../_torch/disaggregation/blob/store.py | 26 ++++++++------ .../blob/test_mooncake_driver.py | 7 ++-- 4 files changed, 37 insertions(+), 35 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/blob/README.md b/tensorrt_llm/_torch/disaggregation/blob/README.md index 975f319a731d..55b33877267a 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/README.md +++ b/tensorrt_llm/_torch/disaggregation/blob/README.md @@ -83,8 +83,9 @@ The boundary is `store.py`; its docstrings are the specification. The rules that name a key twice, and every span lies inside one registered extent, so a driver checks memory once per extent. A driver that cannot take a call refuses it with `DriverRejected`. 4. **Only READY and FAULT are verdicts.** MISS, NOT_READY, DECLINED and CANCELLED say that - nothing was delivered and nothing went wrong, each under a condition the store confirmed; - an error the driver cannot place is FAULT. + nothing was delivered and charge the row no fault, each only under the conditions `RowStatus` + names for it. RPC, transfer and integrity errors are FAULT, and so is any error the driver + cannot place. 5. **`drain` is one-way and never forces an end.** After it the driver accepts no new work; it returns True only when nothing can still touch memory. `close` is allowed only after a True drain and with every registration revoked. diff --git a/tensorrt_llm/_torch/disaggregation/blob/mooncake.py b/tensorrt_llm/_torch/disaggregation/blob/mooncake.py index 22634e404f94..2ff43ae1729b 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/mooncake.py +++ b/tensorrt_llm/_torch/disaggregation/blob/mooncake.py @@ -151,25 +151,20 @@ class _Code(IntEnum): _NO_WORKER_LEFT = faulted("no worker thread is left to run the row") _instance_ids = itertools.count() # names each driver's worker threads -# Why a row was skipped (nothing moved and nothing went wrong), and the result it gets; ``stats()`` -# counts skipped rows by both, and a part's warning names its first row that faulted or was -# skipped, with the skip's reason. -# -# NOT_READY here departs from the protocol's "only for a hold the store confirmed": 0.3.13 cannot -# tell a write in progress from a complete replica it cannot read. -703 answers both, and so does -# batch_is_exist's 0 after a -705. A session left open after a start that raised may not exist, -# and once the master frees one after put_start_release_timeout_sec, its key stays NOT_READY for -# this driver's life, as it is never put again. The usual cause is a write in progress, which as -# a FAULT would stop every fetch that races a writer and count toward isolation; a replica that -# stays unreadable costs only its rows. +# Why a row was skipped (nothing moved, and the row is charged no fault), and the result it gets; +# ``stats()`` counts skipped rows by both, and a part's warning names its first row that faulted +# or was skipped, with the skip's reason. _SKIP_RESULTS = { - # get -703: no complete replica, as while the object is written or once its replica's - # endpoint went stale + # get -703 on a first read: the master holds no readable complete replica, as while the + # object is written or once its replica's endpoint went stale; the slot is untouched "no_complete_replica": _NOT_READY, - # put: a put of this driver is writing the key now, or a session it left open holds the key and - # the master shows no complete object; no session is started + # put: another put of this driver is writing the key now; no session is started "held_by_own_put": _NOT_READY, - "uncommitted": _NOT_READY, # put start -705, and the master shows the object uncommitted + "uncommitted": _NOT_READY, # put start -705, and the master shows no readable complete object + # put: a session this driver left open holds the key and the master shows no complete object. + # The driver never writes such a key again (class docstring), so it refuses the row; the put + # that left the session open already concluded FAULT. + "held_by_session_left_open": _DECLINED, "pool_full": _DECLINED, # put start -200: no space for the object } @@ -302,7 +297,7 @@ class MooncakeDriver: 1. Reserves each key that no put of this driver is writing and no session it left open holds. A key another put is writing is NOT_READY at once. A key whose session it left open is - READY when the master shows its object complete, NOT_READY when it shows none, and FAULT + READY when the master shows its object complete, DECLINED when it shows none, and FAULT when the existence check fails; it is never put again. 2. Starts a session per reserved key, writes every started row in one scatter, ends the session of each row written whole (its object becomes readable) and revokes the rest. @@ -767,8 +762,9 @@ def _put_part(self, rows: _Rows) -> tuple[_PartResults, bool]: closed: set[str] = set() to_master: set[str] = set() try: - if held_by_left_session: # never put again, but another client may have written it - self._confirm_existing(held_by_left_session, results, "held_by_own_put") + # Never put again, but a late end of its own or another client may have completed it. + if held_by_left_session: + self._confirm_existing(held_by_left_session, results, "held_by_session_left_open") if not reserved: return results, True access_ended = self._put_reserved(rows.only(reserved), results, closed, to_master) diff --git a/tensorrt_llm/_torch/disaggregation/blob/store.py b/tensorrt_llm/_torch/disaggregation/blob/store.py index 3f5a50091a43..b99b11decfb2 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/store.py +++ b/tensorrt_llm/_torch/disaggregation/blob/store.py @@ -63,8 +63,8 @@ class RowStatus(enum.Enum): """The logical state of one row within an operation. READY and FAULT are the only two that carry a verdict. MISS, NOT_READY, DECLINED and CANCELLED - say that nothing was delivered and nothing went wrong: the backend reports such a row as not - delivered without a fault, and none of them counts toward isolation. + say that nothing was delivered and charge the row no fault: the backend reports such a row as + not delivered without a fault, and none of them counts toward isolation. """ PENDING = "pending" @@ -82,14 +82,18 @@ class RowStatus(enum.Enum): MISS.""" NOT_READY = "not_ready" - """The key is held by a write that has not finished: a get delivered nothing, a put published - nothing of its own. ``probe`` gives the same answer for such a key. Only for a hold the store - confirmed; a store that cannot tell an unfinished write from another refusal answers FAULT.""" + """A write of the key is known of, but no readable complete object stands under it: a get + delivered nothing and left the row's memory untouched, a put published nothing of its own. + Either the store reported an object of the key without a readable complete replica, as while + a write has not finished or when no complete replica can be read, or another put of the same + driver, still running, writes the key. ``probe`` answers NOT_READY where a get would. RPC, + transfer and integrity errors are FAULT, never NOT_READY.""" DECLINED = "declined" - """For a put: the store would not take the row, for space or policy, and said so. Nothing was - published; the other rows of the operation go on. A refusal the driver cannot place is - FAULT.""" + """For a put: the store would not take the row, for space or policy, and said so; or the + driver will not write the key again because an earlier put of its own left the key unresolved + in the store, and that put concluded FAULT. Nothing was published; the other rows of the + operation go on. A refusal the driver cannot place is FAULT.""" CANCELLED = "cancelled" """Never started, because of a cancel or a drain; no byte moved and the row's memory was not @@ -555,9 +559,9 @@ def probe( Keys the store did not answer in time are UNKNOWN, never MISS; the call returns within the budget whatever the store does, and a driver with no capacity to ask answers UNKNOWN at - once. A key held by an unfinished write is NOT_READY, as a get would report. Called only - when ``capabilities.probe`` is set, and with a hint only when ``capabilities.hints`` is. - Raises on a store error rather than answering MISS. + once. A key a get would find NOT_READY is NOT_READY here too. Called only when + ``capabilities.probe`` is set, and with a hint only when ``capabilities.hints`` is. Raises + on a store error rather than answering MISS. """ ... diff --git a/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py index 4693f19e33fc..6a50529f5fe9 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py +++ b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py @@ -928,7 +928,7 @@ def test_sessions_left_open_are_revoked_only_in_close() -> None: bench.run("get", "C") assert bench.keys_of(PUT_REVOKE) == [key("A")] # only its own first try results, _ = bench.run("put", "A") # the driver holds A's session: it starts no other - assert summary(results) == [("not_ready", None)] # and no complete object exists + assert summary(results) == [("declined", None)] # and no complete object exists assert bench.keys_of(PUT_START).count(key("A")) == 1 assert bench.driver.drain(WAIT_S) is True # sessions left open hold no lender memory @@ -989,7 +989,7 @@ def test_a_put_whose_worker_thread_exits_leaves_its_session_open_for_close( assert key("A") in bench.fake.put_sessions assert bench.driver.stats()["put_sessions_left_for_close"] == 1 results, _ = bench.run("put", "A") # the driver holds A's session: it starts no other - assert summary(results) == [("not_ready", None)] + assert summary(results) == [("declined", None)] assert bench.keys_of(PUT_START).count(key("A")) == 1 wait_for_an_escape(thread_escapes) @@ -1028,7 +1028,8 @@ def test_a_key_whose_session_was_left_open_is_ready_once_its_object_exists(left: bench.run("put", "A") assert bench.driver.stats()[f"put_sessions_left_{left}"] == 1 results, _ = bench.run("put", "A") - assert summary(results) == [("not_ready", None)] # no complete object yet + assert summary(results) == [("declined", None)] # no complete object, and none of its own + assert skipped_counts(bench) == {"rows_declined_held_by_session_left_open": 1} # The master freed the unfinished put, and another client then wrote the key whole. bench.fake.pool.reserved.pop(key("A")) bench.store(A=b"a" * SLOT) From a9f6a74c3cbd44b78e1dd7d776f7ae92979db122 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Thu, 29 Oct 2026 12:42:02 +0000 Subject: [PATCH 09/10] [None][fix] Leave a Mooncake key open when its first put end went unanswered The master matches a session end or revoke to an object by key and client alone, and an end answered RPC_FAIL or RPC_TIMEOUT may still run later, after a retry on another connection ended the session. If the object is then evicted and this driver puts the key again, the late end marks the new, unwritten object complete. Reproduced on a real 0.3.13 master. - A reserved key is now resolved only when the SDK refused its start or all its session calls had a known outcome and left no session. An end whose first answer was RPC_FAIL or RPC_TIMEOUT leaves the key open for close even when the retry ends the session; the row stays READY. A later put of the key is READY while the master shows its object complete, DECLINED once it does not. - A start refused with RPC_FAIL or RPC_TIMEOUT still resolves its key; leaving such keys open would bar every key a master outage met from ever being put again. The class docstring says what a late start can do. - Renames: _sessions_left_open to _keys_left_open, the per-part closed set to resolved, the skip reason held_by_session_left_open to left_open_by_own_put (counter rows_declined_left_open_by_own_put), and the stats put_sessions_left_* to put_keys_left_*. - RowStatus.DECLINED: the driver will not write the key again, as an earlier put of its own may still hold the key or act on a new write of it. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../_torch/disaggregation/blob/mooncake.py | 164 +++++++++--------- .../_torch/disaggregation/blob/store.py | 6 +- .../blob/test_mooncake_driver.py | 66 +++---- 3 files changed, 126 insertions(+), 110 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/blob/mooncake.py b/tensorrt_llm/_torch/disaggregation/blob/mooncake.py index 2ff43ae1729b..fa8702517915 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/mooncake.py +++ b/tensorrt_llm/_torch/disaggregation/blob/mooncake.py @@ -161,16 +161,15 @@ class _Code(IntEnum): # put: another put of this driver is writing the key now; no session is started "held_by_own_put": _NOT_READY, "uncommitted": _NOT_READY, # put start -705, and the master shows no readable complete object - # put: a session this driver left open holds the key and the master shows no complete object. - # The driver never writes such a key again (class docstring), so it refuses the row; the put - # that left the session open already concluded FAULT. - "held_by_session_left_open": _DECLINED, + # put: an earlier put of this driver left the key open (class docstring, step 3) and the + # master shows no complete object. The driver starts no session of such a key again, so it + # refuses the row. + "left_open_by_own_put": _DECLINED, "pool_full": _DECLINED, # put start -200: no space for the object } -# Session-end answers after which the end may still reach the master, even if the retry is refused. -# Such a session is left for close and its key is never put again: the master matches an end by key -# and client only, so a late end would mark complete this client's next object of the key. +# Session-end answers after which the end may still reach the master, so its key is left open +# (class docstring, step 3), even when a retry ends the session. _END_OUTCOME_UNKNOWN = (_Code.RPC_FAIL, _Code.RPC_TIMEOUT, _Code.METADATA_CALL_FAILED) # Session-end answers retried once, as the master ends a session idempotently; an end call that # raised or miscounted is not retried. @@ -178,9 +177,9 @@ class _Code(IntEnum): class _LeftOpen(enum.Enum): - """Who ends a put session this driver left open.""" + """Who frees the session a key left open may still have: close, or the master alone.""" - FOR_CLOSE = "for_close" # close revokes it, as the owner's last step + FOR_CLOSE = "for_close" # close revokes the session the SDK may hold, as the owner's last step TO_MASTER = "to_master" # the target may still write into its object: the master frees it @@ -295,21 +294,32 @@ class MooncakeDriver: A put of a part's rows: - 1. Reserves each key that no put of this driver is writing and no session it left open holds. - A key another put is writing is NOT_READY at once. A key whose session it left open is - READY when the master shows its object complete, DECLINED when it shows none, and FAULT - when the existence check fails; it is never put again. + 1. Reserves each key that no put of this driver is writing or left open. A key another put + is writing is NOT_READY at once. A key a put left open is READY when the master shows its + object complete, DECLINED when it shows none, and FAULT when the existence check fails; + it is never put again. 2. Starts a session per reserved key, writes every started row in one scatter, ends the session of each row written whole (its object becomes readable) and revokes the rest. - 3. Leaves open each session it may have started and did not see ended or revoked: + 3. Resolves each reserved key whose start the SDK refused, or whose session calls all had a + known outcome and left no session; every other key it leaves open for good. The master + matches an end or a revoke to an object by key and client alone, so one that may still + reach it would act on the next put of the key: a late end makes its object readable + before it is written, a late revoke frees its space mid-write. A start refused with + RPC_FAIL or RPC_TIMEOUT may have reached the master, or may still reach it: on a key with + no object it reserves one, so the key's next put is NOT_READY until + ``put_start_discard_timeout_sec``, and it replaces a put of the key unfinished past that + timeout with an unwritten object, which that put's end makes readable when the put is + this driver's. Leaving these keys open would instead bar every key whose start met a + master outage from ever being put again. A key is left open: - to the master, never revoked here, since the target may still write into the object: a scatter that raised or miscounted, a row answered TRANSFER_FAIL. The master frees them after ``put_start_release_timeout_sec``. - for ``close`` to revoke, once no put runs: any other, such as a start that raised or - miscounted, an end whose outcome is unknown (``_END_OUTCOME_UNKNOWN``), a revoke that - failed, or a put stopped by a raise. A put stopped by a raise is unprovable, so no drain - returns True and ``close`` never comes: the master frees its sessions. + miscounted, an end whose outcome is unknown (``_END_OUTCOME_UNKNOWN``) even when a + retried end succeeded, a revoke that failed, or a put stopped by a raise. A put stopped + by a raise is unprovable, so no drain returns True and ``close`` never comes: the + master frees its sessions. Args: store: An SDK store whose ``setup`` returned 0. The driver takes it over: ``close`` @@ -364,10 +374,9 @@ def __init__( # Keys a running put part reserved, until that part ends: no two puts of this driver hold # a session of one key at once, so a -600 is never this driver's own. self._keys_being_put: set[str] = set() - # The put sessions this driver left open (class docstring, step 3), by key. No key is in - # both: a reserved key is left open only as its reservation ends, and is never reserved - # again. - self._sessions_left_open: dict[str, _LeftOpen] = {} + # The keys a put left open (class docstring, step 3). No key is in both: a reserved key is + # left open only as its reservation ends, and is never reserved again. + self._keys_left_open: dict[str, _LeftOpen] = {} self._in_flight_op_count = 0 # operations with rows queued or parts running self._unprovable_op_count = 0 # operations whose access can never be proven ended self._skipped_rows: Counter[str] = Counter() # by reason @@ -479,8 +488,8 @@ def drain(self, budget_s: float) -> bool: """Refuses new work for good, cancels the rows no worker took, ends the worker threads, and waits up to ``budget_s`` for the running parts to end and every worker thread to exit. Returns False at once while an operation is unprovable. Makes no SDK call, so it returns - within ``budget_s``: the sessions left for close hold no lender memory, and ``close`` - revokes them.""" + within ``budget_s``: the sessions of the keys left open for close hold no lender memory, + and ``close`` revokes them.""" with self._lock: self._require_open() self._draining = True @@ -503,11 +512,12 @@ def drain(self, budget_s: float) -> bool: def close(self) -> None: """The owner's last step, once a drain returned True and every region is unregistered: - tries once to revoke the sessions left for close, then closes the SDK client, which takes - no call after that. The revocation waits on the master, and closing holds the GIL through - the unmount, so the owner calls it off the executor loop's thread. A session whose - revocation fails is freed by the master after ``put_start_release_timeout_sec``; sessions - left to the master are never revoked here.""" + tries once to revoke the sessions of the keys left open for close, then closes the SDK + client, which takes no call after that. The SDK revokes the sessions it holds and answers + INVALID_PARAMS for the others. The revocation waits on the master, and closing holds the + GIL through the unmount, so the owner calls it off the executor loop's thread. A session + whose revocation fails is freed by the master after ``put_start_release_timeout_sec``; + the sessions of keys left to the master are never revoked here.""" with self._lock: if self._closed: return @@ -516,7 +526,7 @@ def close(self) -> None: if any(r.live for r in self._registrations): raise DriverRejected("close needs every registration unregistered") self._closed = True - left = self._sessions_left_open.items() + left = self._keys_left_open.items() to_revoke = [key for key, by in left if by is _LeftOpen.FOR_CLOSE] try: if to_revoke: # no put runs any more, so none of them is in use @@ -524,22 +534,22 @@ def close(self) -> None: with self._lock: for key, code in zip(to_revoke, codes): if _session_gone(code): - del self._sessions_left_open[key] + del self._keys_left_open[key] finally: self._store.close() def stats(self) -> dict[str, int]: """A snapshot for logs and tests; not part of the driver protocol.""" with self._lock: - left = list(self._sessions_left_open.values()) + left = list(self._keys_left_open.values()) return { **{ f"rows_{r.status.value}_{reason}": self._skipped_rows[reason] for reason, r in _SKIP_RESULTS.items() }, "ops_unprovable": self._unprovable_op_count, - "put_sessions_left_for_close": left.count(_LeftOpen.FOR_CLOSE), - "put_sessions_left_to_master": left.count(_LeftOpen.TO_MASTER), + "put_keys_left_for_close": left.count(_LeftOpen.FOR_CLOSE), + "put_keys_left_to_master": left.count(_LeftOpen.TO_MASTER), } # Admission, on the caller's thread. @@ -753,68 +763,65 @@ def _call_on_change(self, op: _Op) -> None: def _put_part(self, rows: _Rows) -> tuple[_PartResults, bool]: """Returns the part's results and whether access ended. Whatever stops the calls, each - reserved key whose session they did not see closed is left open (class docstring, - step 3).""" + reserved key they did not resolve is left open (class docstring, step 3).""" results = _PartResults() - reserved, held_by_left_session = self._reserve_keys(rows.keys, results) - # Reserved keys known to hold no session (never opened, or ended or revoked), and those - # whose object the target may still write into. - closed: set[str] = set() + reserved, already_left_open = self._reserve_keys(rows.keys, results) + # Reserved keys a later put may reserve again (class docstring, step 3), and those whose + # object the target may still write into. + resolved: set[str] = set() to_master: set[str] = set() try: - # Never put again, but a late end of its own or another client may have completed it. - if held_by_left_session: - self._confirm_existing(held_by_left_session, results, "held_by_session_left_open") + # Never put again: READY only while the master shows its object complete. + if already_left_open: + self._confirm_existing(already_left_open, results, "left_open_by_own_put") if not reserved: return results, True - access_ended = self._put_reserved(rows.only(reserved), results, closed, to_master) + access_ended = self._put_reserved(rows.only(reserved), results, resolved, to_master) return results, access_ended finally: - self._release_keys_and_record_left_open(reserved, closed, to_master) + self._release_keys_and_record_left_open(reserved, resolved, to_master) def _reserve_keys(self, keys: list[str], results: _PartResults) -> tuple[list[str], list[str]]: - """Reserves every key that no put of this driver is writing and no session it left open - holds, and returns them, then the keys a left-open session holds. A key another put of - this driver is writing is NOT_READY.""" + """Reserves every key that no put of this driver is writing or left open, and returns + them, then the keys left open. A key another put of this driver is writing is + NOT_READY.""" with self._lock: - being_put, left_open = self._keys_being_put, self._sessions_left_open.keys() + being_put, left_open = self._keys_being_put, self._keys_left_open.keys() if being_put.isdisjoint(keys) and left_open.isdisjoint(keys): # usually - reserved, held_by_left_session = keys, [] + reserved, already_left_open = keys, [] else: - reserved, held_by_left_session = [], [] + reserved, already_left_open = [], [] for key in keys: if key in being_put: results.skip(key, "held_by_own_put") elif key in left_open: - held_by_left_session.append(key) + already_left_open.append(key) else: reserved.append(key) self._keys_being_put.update(reserved) - return reserved, held_by_left_session + return reserved, already_left_open def _release_keys_and_record_left_open( - self, reserved: list[str], closed: set[str], to_master: set[str] + self, reserved: list[str], resolved: set[str], to_master: set[str] ) -> None: - """Ends the reservation of ``reserved`` and records each session not ``closed`` as left - open: for the master when its key is in ``to_master``, otherwise for ``close``.""" + """Ends the reservation of ``reserved`` and records each key not ``resolved`` as left + open: for the master when it is in ``to_master``, otherwise for ``close``.""" left: dict[str, _LeftOpen] = {} - if not closed.issuperset(reserved): # usually every session ended + if not resolved.issuperset(reserved): # usually every key resolved left = { key: _LeftOpen.TO_MASTER if key in to_master else _LeftOpen.FOR_CLOSE for key in reserved - if key not in closed + if key not in resolved } with self._lock: self._keys_being_put.difference_update(reserved) - self._sessions_left_open.update(left) + self._keys_left_open.update(left) if left: first, why = next(iter(left.items())) - logger.warning( - f"mooncake: {len(left)} put sessions left open, first {first} {why.value}" - ) + logger.warning(f"mooncake: {len(left)} put keys left open, first {first} {why.value}") def _put_reserved( - self, rows: _Rows, results: _PartResults, closed: set[str], to_master: set[str] + self, rows: _Rows, results: _PartResults, resolved: set[str], to_master: set[str] ) -> bool: """Starts a session per row, writes and finishes the started ones, and confirms the rows whose object exists already. Returns whether access ended.""" @@ -826,8 +833,8 @@ def _put_reserved( if code == _Code.OK: started.append(key) continue - if code != _Code.METADATA_CALL_FAILED: # the SDK answered: no session opened - closed.add(key) + if code != _Code.METADATA_CALL_FAILED: # the SDK answered and holds no session + resolved.add(key) if code == _Code.OBJECT_ALREADY_EXISTS: existing.append(key) elif code == _Code.NO_AVAILABLE_HANDLE: @@ -840,13 +847,13 @@ def _put_reserved( results.fault([key], f"put session start: {_describe(code)}") access_ended = True if started: - access_ended = self._write_and_finish(rows.only(started), results, closed, to_master) + access_ended = self._write_and_finish(rows.only(started), results, resolved, to_master) if existing: # readable only once the master shows it committed self._confirm_existing(existing, results, "uncommitted") return access_ended def _write_and_finish( - self, rows: _Rows, results: _PartResults, closed: set[str], to_master: set[str] + self, rows: _Rows, results: _PartResults, resolved: set[str], to_master: set[str] ) -> bool: """Writes the rows whose put session started in one scatter, then ends the session of each row written whole and revokes the rest, except those the target may still write into. @@ -884,20 +891,22 @@ def _write_and_finish( else: to_revoke.append(key) if written: - to_revoke += self._end_sessions(written, results, closed) + to_revoke += self._end_sessions(written, results, resolved) if to_revoke: codes = self._metadata_call(self._store.batch_put_session_revoke, to_revoke) - closed.update(key for key, code in zip(to_revoke, codes) if _session_gone(code)) + resolved.update(key for key, code in zip(to_revoke, codes) if _session_gone(code)) return True - def _end_sessions(self, keys: list[str], results: _PartResults, closed: set[str]) -> list[str]: + def _end_sessions( + self, keys: list[str], results: _PartResults, resolved: set[str] + ) -> list[str]: """Ends the sessions of rows written whole. Returns the keys whose end was refused, to - revoke; a session whose end may still reach the master, answered ``_END_OUTCOME_UNKNOWN`` - first, stays open even when the retry was refused.""" + revoke. A key whose first end answered ``_END_OUTCOME_UNKNOWN`` stays open, as that end may + still reach the master; its row is READY when the retry ended the session.""" end = self._store.batch_put_session_end codes = self._metadata_call(end, keys) if not any(codes): # every session ended (OK is 0) - closed.update(keys) + resolved.update(keys) return [] final = dict(zip(keys, codes)) retry = [key for key, code in final.items() if code in _END_RETRIED] @@ -905,12 +914,11 @@ def _end_sessions(self, keys: list[str], results: _PartResults, closed: set[str] final.update(zip(retry, self._metadata_call(end, retry))) refused = [] for key, first in zip(keys, codes): - code = final[key] - if code == _Code.OK: - closed.add(key) - continue - results.fault([key], f"put session end: {_describe(code)}") - if first not in _END_OUTCOME_UNKNOWN: + if final[key] != _Code.OK: + results.fault([key], f"put session end: {_describe(final[key])}") + if first == _Code.OK: + resolved.add(key) + elif first not in _END_OUTCOME_UNKNOWN: refused.append(key) return refused diff --git a/tensorrt_llm/_torch/disaggregation/blob/store.py b/tensorrt_llm/_torch/disaggregation/blob/store.py index b99b11decfb2..b3b132ec02dc 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/store.py +++ b/tensorrt_llm/_torch/disaggregation/blob/store.py @@ -91,9 +91,9 @@ class RowStatus(enum.Enum): DECLINED = "declined" """For a put: the store would not take the row, for space or policy, and said so; or the - driver will not write the key again because an earlier put of its own left the key unresolved - in the store, and that put concluded FAULT. Nothing was published; the other rows of the - operation go on. A refusal the driver cannot place is FAULT.""" + driver will not write the key again, as an earlier put of its own may still hold the key or + act on a new write of it. Nothing was published; the other rows of the operation go on. A + refusal the driver cannot place is FAULT.""" CANCELLED = "cancelled" """Never started, because of a cancel or a drain; no byte moved and the row's memory was not diff --git a/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py index 6a50529f5fe9..63193ad83169 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py +++ b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py @@ -734,7 +734,10 @@ def test_get_answers_map_to_row_statuses(case: str) -> None: # master's object cap. "INVALID_PARAMS": (-600, "fault", "put session start: INVALID_PARAMS(-600)", None), "INVALID_REPLICA": (-702, "fault", "put session start: INVALID_REPLICA(-702)", None), + # The start may still reach the master, yet its key is not left open (MooncakeDriver + # docstring, step 3). "RPC_FAIL": (-900, "fault", "put session start: RPC_FAIL(-900)", None), + "RPC_TIMEOUT": (-901, "fault", "put session start: RPC_TIMEOUT(-901)", None), "unknown_code": (-9999, "fault", "put session start: code(-9999)", None), } @@ -752,6 +755,7 @@ def test_put_session_start_answers_map_to_row_statuses(case: str) -> None: assert len(start.args) == 2 and start.kwargs == {} # never a replicate config for method in (PUT_RANGES, PUT_END, PUT_REVOKE, "batch_is_exist"): assert bench.fake.calls_to(method) == [] # a session it did not start is not its own + assert bench.driver.stats()["put_keys_left_for_close"] == 0 PUT_WRITE_ANSWERS = { @@ -783,8 +787,8 @@ def test_put_ranges_and_end_answers_map_to_row_statuses(case: str) -> None: assert (bench.fake.pool.objects.get(key("A")) == b"aaaaaaaa") == (status == "ready") assert bench.keys_of(PUT_REVOKE) == ([key("A")] if session == "revoked" else []) stats = bench.driver.stats() - assert stats["put_sessions_left_for_close"] == (session == "left_for_close") - assert stats["put_sessions_left_to_master"] == (session == "left_to_master") + assert stats["put_keys_left_for_close"] == (session == "left_for_close") + assert stats["put_keys_left_to_master"] == (session == "left_to_master") EXISTENCE_ANSWERS = { @@ -824,8 +828,8 @@ class SessionScenario(NamedTuple): rows: str # each row ready (R) or faulted (F) access_ends: bool = True revoked: str = "" # rows whose session the put revokes, or tries to - left_for_close: str = "" # rows whose session is left open for close to revoke - left_to_master: str = "" # rows whose session is left open for the master alone to free + left_for_close: str = "" # rows whose key is left open for close + left_to_master: str = "" # rows whose key is left open for the master alone to free SESSION_SCENARIOS = { @@ -833,7 +837,10 @@ class SessionScenario(NamedTuple): # A TCP target may still write into B's object space: only the master may free it. "transfer_fail": SessionScenario({PUT_RANGES: [[None, -800, None]]}, "RFR", left_to_master="B"), "short_write": SessionScenario({PUT_RANGES: [[None, 5, None]]}, "RFR", revoked="B"), - "end_retried": SessionScenario({PUT_END: [[-900, None, None], [None]]}, "RRR"), + # The retry ends A's session, but the first end may still reach the master. + "end_retried": SessionScenario( + {PUT_END: [[-900, None, None], [None]]}, "RRR", left_for_close="A" + ), # Ends that may still reach the master. "end_fails_twice": SessionScenario( {PUT_END: [[-901, None, None], [-900]]}, "FRR", left_for_close="A" @@ -892,8 +899,8 @@ def test_every_started_put_session_is_ended_revoked_or_left_open(name: str) -> N for_close = {key(n) for n in scenario.left_for_close} to_master = {key(n) for n in scenario.left_to_master} stats = bench.driver.stats() - assert stats["put_sessions_left_for_close"] == len(for_close) - assert stats["put_sessions_left_to_master"] == len(to_master) + assert stats["put_keys_left_for_close"] == len(for_close) + assert stats["put_keys_left_to_master"] == len(to_master) # Every session the SDK still holds is left open; the SDK may have dropped one the driver # could not see end, as when an end landed but its reply did not. assert bench.fake.put_sessions <= for_close | to_master @@ -902,15 +909,16 @@ def test_every_started_put_session_is_ended_revoked_or_left_open(name: str) -> N assert set(bench.keys_of(PUT_END)) <= set(written) ready = [key(n) for n, r in zip("ABC", results) if r.status is RowStatus.READY] assert set(ready) == set(bench.fake.keys_answered(PUT_END, 0)) - # Close revokes the sessions left for it, never those left to the master; drain revokes none. + # Close revokes the sessions of the keys left for it, never those left to the master; drain + # revokes none. if scenario.access_ends: revocations = len(bench.keys_of(PUT_REVOKE)) assert bench.driver.drain(WAIT_S) is True assert len(bench.keys_of(PUT_REVOKE)) == revocations bench.driver.unregister(bench.handle) bench.driver.close() - assert bench.driver.stats()["put_sessions_left_for_close"] == 0 - assert bench.driver.stats()["put_sessions_left_to_master"] == len(to_master) + assert bench.driver.stats()["put_keys_left_for_close"] == 0 + assert bench.driver.stats()["put_keys_left_to_master"] == len(to_master) assert bench.fake.put_sessions == to_master assert not to_master & set(bench.keys_of(PUT_REVOKE)) @@ -921,7 +929,7 @@ def test_sessions_left_open_are_revoked_only_in_close() -> None: bench.fake.script(PUT_END, [-704]) bench.fake.script(PUT_REVOKE, [-900]) bench.run("put", "A") - assert bench.driver.stats()["put_sessions_left_for_close"] == 1 + assert bench.driver.stats()["put_keys_left_for_close"] == 1 bench.write(1, b"bbbbbbbb") bench.store(C=b"cccccccc") bench.run("put", "X", "B") @@ -931,13 +939,13 @@ def test_sessions_left_open_are_revoked_only_in_close() -> None: assert summary(results) == [("declined", None)] # and no complete object exists assert bench.keys_of(PUT_START).count(key("A")) == 1 - assert bench.driver.drain(WAIT_S) is True # sessions left open hold no lender memory + assert bench.driver.drain(WAIT_S) is True # sessions of keys left open hold no lender memory assert bench.keys_of(PUT_REVOKE) == [key("A")] bench.driver.unregister(bench.handle) bench.fake.script(PUT_REVOKE, [-900]) bench.driver.close() assert bench.keys_of(PUT_REVOKE) == [key("A")] * 2 - assert bench.driver.stats()["put_sessions_left_for_close"] == 1 # the master frees it later + assert bench.driver.stats()["put_keys_left_for_close"] == 1 # the master frees it later assert len(bench.fake.calls_to("close")) == 1 @@ -947,7 +955,7 @@ def test_drain_returns_at_once_while_a_revocation_would_block() -> None: bench.fake.script(PUT_END, [-901]) bench.fake.script(PUT_END, [-901]) # the end and its retry go unanswered bench.run("put", "A") - assert bench.driver.stats()["put_sessions_left_for_close"] == 1 + assert bench.driver.stats()["put_keys_left_for_close"] == 1 gate = bench.fake.block(PUT_REVOKE) started = time.monotonic() assert bench.driver.drain(WAIT_S) is True # once the idle workers exit @@ -957,7 +965,7 @@ def test_drain_returns_at_once_while_a_revocation_would_block() -> None: bench.driver.unregister(bench.handle) bench.driver.close() assert bench.keys_of(PUT_REVOKE) == [key("A")] - assert bench.driver.stats()["put_sessions_left_for_close"] == 0 + assert bench.driver.stats()["put_keys_left_for_close"] == 0 assert bench.fake.put_sessions == set() @@ -987,7 +995,7 @@ def test_a_put_whose_worker_thread_exits_leaves_its_session_open_for_close( assert summary(results) == [("fault", "the worker stopped before the row was decided")] assert not access_ended assert key("A") in bench.fake.put_sessions - assert bench.driver.stats()["put_sessions_left_for_close"] == 1 + assert bench.driver.stats()["put_keys_left_for_close"] == 1 results, _ = bench.run("put", "A") # the driver holds A's session: it starts no other assert summary(results) == [("declined", None)] assert bench.keys_of(PUT_START).count(key("A")) == 1 @@ -1016,26 +1024,26 @@ def test_a_key_another_put_of_this_driver_holds_is_not_ready_without_asking_the_ assert summary(results) == [("ready", None)] # it exists already, committed -@pytest.mark.parametrize("left", ["for_close", "to_master"]) -def test_a_key_whose_session_was_left_open_is_ready_once_its_object_exists(left: str) -> None: +@pytest.mark.parametrize("left", ["to_master", "for_close"]) +def test_a_key_left_open_is_never_put_again_and_ready_once_its_object_exists(left: str) -> None: bench = Bench() bench.fill_sources("put", "A") if left == "to_master": bench.fake.script(PUT_RANGES, [-800]) - else: # the end and its retry go unanswered - bench.fake.script(PUT_END, [-901]) + else: # the retry ends the session, but the first end may still land bench.fake.script(PUT_END, [-901]) bench.run("put", "A") - assert bench.driver.stats()[f"put_sessions_left_{left}"] == 1 + assert bench.driver.stats()[f"put_keys_left_{left}"] == 1 + # The store dropped what the put left there: the master freed it, or evicted the object. + bench.fake.pool.reserved.pop(key("A"), None) + bench.fake.pool.objects.pop(key("A"), None) results, _ = bench.run("put", "A") - assert summary(results) == [("declined", None)] # no complete object, and none of its own - assert skipped_counts(bench) == {"rows_declined_held_by_session_left_open": 1} - # The master freed the unfinished put, and another client then wrote the key whole. - bench.fake.pool.reserved.pop(key("A")) - bench.store(A=b"a" * SLOT) + assert summary(results) == [("declined", None)] + assert skipped_counts(bench) == {"rows_declined_left_open_by_own_put": 1} + bench.store(A=b"a" * SLOT) # another client then wrote the key whole results, _ = bench.run("put", "A") assert summary(results) == [("ready", None)] # whoever wrote it - assert bench.keys_of(PUT_START) == [key("A")] # its session was never started again + assert bench.keys_of(PUT_START) == [key("A")] # no session of it was started again assert bench.keys_of("batch_is_exist") == [key("A")] * 2 @@ -1066,7 +1074,7 @@ def test_an_unprovable_put_leaves_its_session_to_the_master() -> None: assert not access_ended assert bench.driver.drain(0.0) is False assert bench.fake.calls_to(PUT_REVOKE) == [] - assert bench.driver.stats()["put_sessions_left_to_master"] == 1 + assert bench.driver.stats()["put_keys_left_to_master"] == 1 LEASE_EXPIRED_CASES = { @@ -1165,7 +1173,7 @@ def broken(*_): assert summary(results) == [("fault", f"{direction} failed: ValueError('mapping bug')")] assert not access_ended # A put's session may still be open: the driver keeps holding the key. - assert bench.driver.stats()["put_sessions_left_for_close"] == (direction == "put") + assert bench.driver.stats()["put_keys_left_for_close"] == (direction == "put") # Admission, registration, concurrency. From 673a274729ac6f09b70e7d94517713e973564b4f Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Fri, 30 Oct 2026 08:16:53 +0000 Subject: [PATCH 10/10] [None][test] Pin that a Mooncake put holds its slot until its scatter returns The Mooncake scatter reads the lender's slots directly, so a put whose access reads as ended while the scatter still runs lets the lender reuse memory the store is still reading. Until now only the conformance suite's blocked-put check caught that. #21314 drops that assertion from the suite, since a driver that stages the payload in memory of its own may truthfully end access early. This driver test keeps the property pinned for Mooncake: with the scatter held inside the SDK, the put's access has not ended, and once it returns the row is READY and access has ended. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../disaggregation/blob/test_mooncake_driver.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py index 63193ad83169..71d7eaa03e4c 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py +++ b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py @@ -409,6 +409,20 @@ def test_a_cancel_ends_only_the_rows_no_worker_took() -> None: assert bench.driver.drain(WAIT_S) is True +def test_a_put_holds_its_slot_until_its_scatter_returns() -> None: + bench = Bench() + bench.fill_sources("put", "A") + gate = bench.fake.block(PUT_RANGES) + op = bench.driver.put(bench.rows("A")) + try: + assert gate.entered.wait(WAIT_S) + assert not bench.driver.access_ended(op) # the scatter reads the slot itself + finally: + gate.release.set() + results, access_ended = bench.settle(op) + assert summary(results) == [("ready", None)] and access_ended + + def test_drain_cancels_the_rows_no_worker_took() -> None: bench = Bench(workers=1, bytes_per_call=SLOT) bench.store(A=b"aaaaaaaa", B=b"bbbbbbbb")