diff --git a/tensorrt_llm/_torch/disaggregation/blob/README.md b/tensorrt_llm/_torch/disaggregation/blob/README.md index 223f1a0bc9c9..55b33877267a 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/README.md +++ b/tensorrt_llm/_torch/disaggregation/blob/README.md @@ -32,6 +32,8 @@ 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 | +| `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 @@ -40,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 @@ -52,6 +54,12 @@ 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. +- `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 @@ -75,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/__init__.py b/tensorrt_llm/_torch/disaggregation/blob/__init__.py index 80de5d3d1a62..56e778597f47 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/blob/__init__.py @@ -12,7 +12,13 @@ - ``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. +- ``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.py b/tensorrt_llm/_torch/disaggregation/blob/mooncake.py new file mode 100644 index 000000000000..fa8702517915 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/blob/mooncake.py @@ -0,0 +1,1060 @@ +# 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 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 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: 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: 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, 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. +_END_RETRIED = (_Code.RPC_FAIL, _Code.RPC_TIMEOUT) + + +class _LeftOpen(enum.Enum): + """Who frees the session a key left open may still have: close, or the master alone.""" + + 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 + + +@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 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. 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``) 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`` + 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. + 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 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 + # 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. 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) + 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. Makes no SDK call, so it returns + 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 + 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 + return drained + + 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 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 + 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 + 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 + 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._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._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_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. + + 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 + # 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) + + 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 they did not resolve is left open (class docstring, step 3).""" + results = _PartResults() + 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: 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, resolved, to_master) + return results, access_ended + finally: + 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 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._keys_left_open.keys() + if being_put.isdisjoint(keys) and left_open.isdisjoint(keys): # usually + reserved, already_left_open = keys, [] + else: + 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: + already_left_open.append(key) + else: + reserved.append(key) + self._keys_being_put.update(reserved) + return reserved, already_left_open + + def _release_keys_and_record_left_open( + self, reserved: list[str], resolved: set[str], to_master: set[str] + ) -> None: + """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 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 resolved + } + with self._lock: + self._keys_being_put.difference_update(reserved) + self._keys_left_open.update(left) + if left: + first, why = next(iter(left.items())) + logger.warning(f"mooncake: {len(left)} put keys left open, first {first} {why.value}") + + def _put_reserved( + 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.""" + 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 and holds no session + resolved.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, 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, 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. + 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, resolved) + if to_revoke: + codes = self._metadata_call(self._store.batch_put_session_revoke, to_revoke) + 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, resolved: set[str] + ) -> list[str]: + """Ends the sessions of rows written whole. Returns the keys whose end was refused, to + 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) + resolved.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): + 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 + + 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/tensorrt_llm/_torch/disaggregation/blob/mooncake_setup.py b/tensorrt_llm/_torch/disaggregation/blob/mooncake_setup.py new file mode 100644 index 000000000000..5683fa30cd8d --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/blob/mooncake_setup.py @@ -0,0 +1,415 @@ +# 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. + + 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 + ``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/tensorrt_llm/_torch/disaggregation/blob/store.py b/tensorrt_llm/_torch/disaggregation/blob/store.py index 3f5a50091a43..b3b132ec02dc 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, 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 @@ -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/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/fake_mooncake.py b/tests/unittest/_torch/disaggregation/blob/fake_mooncake.py new file mode 100644 index 000000000000..53771b65b1d6 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/fake_mooncake.py @@ -0,0 +1,303 @@ +# 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.closed = False + 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", (), {}) + self.closed = True + 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) + 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 + 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..71d7eaa03e4c --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py @@ -0,0 +1,1518 @@ +# 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_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") + 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), + # 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), +} + + +@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 + assert bench.driver.stats()["put_keys_left_for_close"] == 0 + + +PUT_WRITE_ANSWERS = { + # case: (the scatter's answer, the end's answer (None: real), status, fault, the session's end: + # 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. + "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_close"), +} + + +@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_keys_left_for_close"] == (session == "left_for_close") + assert stats["put_keys_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_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 = { + "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"), + # 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" + ), + "end_retry_refused": SessionScenario( + {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_close="A" + ), + "end_raises": SessionScenario( + {PUT_END: [RuntimeError("rpc lost")]}, "FFF", left_for_close="ABC" + ), + "end_committed_then_raised": SessionScenario( + {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_close="A", + ), + # The SDK may have opened any of the sessions. + "start_raises": SessionScenario( + {PUT_START: [RuntimeError("rpc lost")]}, "FFF", left_for_close="ABC" + ), + "start_miscounts": SessionScenario( + {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( + {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_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_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 + # 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)) + # 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_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)) + + +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_keys_left_for_close"] == 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) == [("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 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_keys_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_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 + 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_keys_left_for_close"] == 0 + 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: + 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_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 + 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", ["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 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_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)] + 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")] # no session of it was 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_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") + assert not access_ended + assert bench.driver.drain(0.0) is False + assert bench.fake.calls_to(PUT_REVOKE) == [] + assert bench.driver.stats()["put_keys_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_keys_left_for_close"] == (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.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 + 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() 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..01707545b1b2 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_mooncake_real.py @@ -0,0 +1,451 @@ +# 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. + A closed endpoint's driver has closed its client already, and closing a client twice does + nothing. + """ + + 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 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..ca4e47eede9a --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_mooncake_setup.py @@ -0,0 +1,575 @@ +# 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 + 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 == { + "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 + driver.close() + 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 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..1cff9485d5a5 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py @@ -0,0 +1,266 @@ +# 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 contextlib +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.blob.store import DriverRejected +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. 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._drivers = [] + self.backends = [] + + def open_driver(self): + """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) + 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.""" + 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 driver in self._drivers: + with contextlib.suppress(DriverRejected): # one a failed build never drained + driver.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, 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 + ] + 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() + binding.driver.close() + 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() + binding.driver.close() + + +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() + binding.driver.close()