[None][feat] Add a Mooncake store driver for the blob backend

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.

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.
 """

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.
         """
         ...
 

tests/test_common/mooncake_utils.py

@@ -12,28 +12,52 @@
 # 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
 other ports belong to whatever else runs there. A port is reserved here by
 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()