From 43d062d181ac035b3380b7638cf0b8ef7b087af2 Mon Sep 17 00:00:00 2001 From: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:38:24 +0800 Subject: [PATCH 1/6] [None][feat] Add the blob store driver boundary and an in-memory test driver Defines the private StoreDriver contract, object format and key encoding for KV rows shared through a byte store. An in-memory store and a controlled driver let the tests pin the contract on CPU. Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> --- .../_torch/disaggregation/blob/__init__.py | 7 + .../_torch/disaggregation/blob/format.py | 49 ++ .../_torch/disaggregation/blob/keys.py | 28 + .../_torch/disaggregation/blob/store.py | 256 ++++++++ .../_torch/disaggregation/blob/host_memory.py | 50 ++ .../disaggregation/blob/memory_store.py | 373 ++++++++++++ .../disaggregation/blob/test_memory_store.py | 566 ++++++++++++++++++ 7 files changed, 1329 insertions(+) create mode 100644 tensorrt_llm/_torch/disaggregation/blob/__init__.py create mode 100644 tensorrt_llm/_torch/disaggregation/blob/format.py create mode 100644 tensorrt_llm/_torch/disaggregation/blob/keys.py create mode 100644 tensorrt_llm/_torch/disaggregation/blob/store.py create mode 100644 tests/unittest/_torch/disaggregation/blob/host_memory.py create mode 100644 tests/unittest/_torch/disaggregation/blob/memory_store.py create mode 100644 tests/unittest/_torch/disaggregation/blob/test_memory_store.py diff --git a/tensorrt_llm/_torch/disaggregation/blob/__init__.py b/tensorrt_llm/_torch/disaggregation/blob/__init__.py new file mode 100644 index 000000000000..7c7756c302ba --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/blob/__init__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The blob backend: KV rows shared through a byte store. + +Only the private driver boundary lives here so far. Nothing in this package imports the public +backend contract or an optional store SDK at import time. +""" diff --git a/tensorrt_llm/_torch/disaggregation/blob/format.py b/tensorrt_llm/_torch/disaggregation/blob/format.py new file mode 100644 index 000000000000..9875fb8227f8 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/blob/format.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The object format: the payload bytes followed by a fixed-size trailer. + +Private and versioned. The trailer lets a reader check that an object has the expected format, +belongs to the expected name and has the expected length; it is not a payload checksum and proves +neither a complete commit nor the end of access. Any incompatible change bumps ``FORMAT_VERSION``. +""" + +from __future__ import annotations + +import hashlib +import struct + +FORMAT_VERSION = 1 +MAGIC = b"TLBK" + +# magic, version, 2 bytes of padding, name digest, payload length. +_TRAILER = struct.Struct("<4sH2x16sQ") +TRAILER_BYTES = _TRAILER.size +_DIGEST_BYTES = 16 + + +def name_digest(name: bytes) -> bytes: + """The 16-byte digest of an object name that the trailer carries.""" + return hashlib.blake2b(name, digest_size=_DIGEST_BYTES).digest() + + +def encode_trailer(name: bytes, payload_nbytes: int) -> bytes: + """The trailer for an object of ``payload_nbytes`` bytes stored under ``name``.""" + if payload_nbytes <= 0: + raise ValueError(f"payload length must be positive, got {payload_nbytes}") + return _TRAILER.pack(MAGIC, FORMAT_VERSION, name_digest(name), payload_nbytes) + + +def check_trailer(trailer: bytes, name: bytes, payload_nbytes: int) -> str | None: + """``None`` when ``trailer`` matches ``name`` and ``payload_nbytes``, else why it does not.""" + if len(trailer) != TRAILER_BYTES: + return f"trailer is {len(trailer)} bytes, expected {TRAILER_BYTES}" + magic, version, digest, length = _TRAILER.unpack(trailer) + if magic != MAGIC: + return "bad magic" + if version != FORMAT_VERSION: + return f"format version {version}, expected {FORMAT_VERSION}" + if digest != name_digest(name): + return "object name mismatch" + if length != payload_nbytes: + return f"payload length {length}, expected {payload_nbytes}" + return None diff --git a/tensorrt_llm/_torch/disaggregation/blob/keys.py b/tensorrt_llm/_torch/disaggregation/blob/keys.py new file mode 100644 index 000000000000..c8fbc39f8722 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/blob/keys.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Store keys for blob objects. + +A private, versioned encoding of the lender's opaque object names, not a stable public format. An +incompatible change bumps ``KEY_FORMAT_VERSION`` so objects written under the old one miss instead +of being misread. +""" + +from __future__ import annotations + +KEY_FORMAT_VERSION = 1 + +_SEPARATOR = "/" + + +def encode_key(prefix: str, name: bytes) -> str: + """``"/v/"``. + + Args: + prefix: The deployment's namespace; non-empty, without ``/`` or whitespace. + name: The lender's opaque object name; encoded whole, never parsed. + """ + if not prefix or _SEPARATOR in prefix or any(c.isspace() for c in prefix): + raise ValueError(f"bad key prefix {prefix!r}: non-empty, no '/' and no whitespace") + if not name: + raise ValueError("an object name is non-empty") + return f"{prefix}{_SEPARATOR}v{KEY_FORMAT_VERSION}{_SEPARATOR}{name.hex()}" diff --git a/tensorrt_llm/_torch/disaggregation/blob/store.py b/tensorrt_llm/_torch/disaggregation/blob/store.py new file mode 100644 index 000000000000..eac8c10a9264 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/blob/store.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The private boundary between the blob backend and a byte store. + +A driver moves whole rows between registered host memory and one store. It answers two questions +separately: what each row's logical result is (:meth:`StoreDriver.poll`), and whether an operation +can still touch the caller's memory (:meth:`StoreDriver.access_ended`). Neither answer implies the +other. Nothing here belongs to the public backend contract in ``base/``; the backend translates +between the two. +""" + +from __future__ import annotations + +import enum +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Protocol, runtime_checkable + + +class StoreError(Exception): + """A store-side failure the driver could not express as a row result.""" + + +class RegistrationError(StoreError): + """Registering a region failed. + + Args: + reason: What went wrong. + handle: The handle to ``unregister`` later when the store may have registered the region + anyway, or ``None`` when nothing escaped. + """ + + def __init__(self, reason: str, handle: RegionHandle | None = None) -> None: + super().__init__(reason) + self.handle = handle + + +class RowStatus(enum.Enum): + """The logical state of one row within an operation.""" + + PENDING = "pending" + """Not concluded yet.""" + + READY = "ready" + """For a get: the bytes landed whole and the driver will not write the row again. For a put: + the object is readable whole under its key.""" + + MISS = "miss" + """For a get: the store holds no object under the key. Not a failure.""" + + CANCELLED = "cancelled" + """Never started, because of a cancel or a drain; the row's memory was not touched.""" + + FAULT = "fault" + """Anything else: timeout, size or format mismatch, transport or store error.""" + + +@dataclass(frozen=True) +class RowResult: + """One row's logical result. ``fault`` is set exactly when ``status`` is ``FAULT``.""" + + status: RowStatus + fault: str | None = None + + def __post_init__(self) -> None: + if (self.status is RowStatus.FAULT) != (self.fault is not None): + raise ValueError("a fault text is given exactly for FAULT results") + + @property + def concluded(self) -> bool: + return self.status is not RowStatus.PENDING + + +PENDING = RowResult(RowStatus.PENDING) +READY = RowResult(RowStatus.READY) +MISS = RowResult(RowStatus.MISS) +CANCELLED = RowResult(RowStatus.CANCELLED) + + +def faulted(reason: str) -> RowResult: + """A FAULT result carrying ``reason``.""" + return RowResult(RowStatus.FAULT, reason) + + +def concluded(results: Sequence[RowResult]) -> bool: + """Whether no row of ``results`` is pending.""" + return all(result.concluded for result in results) + + +@dataclass(frozen=True) +class Span: + """Host memory ``[address, address + nbytes)``.""" + + address: int + nbytes: int + + def __post_init__(self) -> None: + if self.address < 0 or self.nbytes <= 0: + raise ValueError(f"bad span: address={self.address}, nbytes={self.nbytes}") + + @property + def end(self) -> int: + return self.address + self.nbytes + + def overlaps(self, other: Span) -> bool: + return self.address < other.end and other.address < self.end + + def within(self, other: Span) -> bool: + return other.address <= self.address and self.end <= other.end + + +@dataclass(frozen=True) +class RegionLayout: + """How a registered region is divided: ``slot_count`` slots of ``slot_bytes``, back to back. + + ``label`` is an opaque layout label, equal for regions laid out alike; it is not an identity. + """ + + slot_bytes: int + slot_count: int + label: str = "" + + def __post_init__(self) -> None: + if self.slot_bytes <= 0 or self.slot_count <= 0: + raise ValueError( + f"bad layout: slot_bytes={self.slot_bytes}, slot_count={self.slot_count}" + ) + + @property + def nbytes(self) -> int: + return self.slot_bytes * self.slot_count + + +@dataclass(frozen=True) +class RowIO: + """One row to move: its store key, where its payload sits and where its trailer sits. + + The payload is the KV bytes of one staging slot; the trailer is the format's fixed tail, in a + buffer of the backend's own. The two spans never overlap. ``position`` is ``(group, row)`` + within the submission, kept so the backend can map results back; drivers treat it as opaque. + """ + + key: str + payload: Span + trailer: Span + position: tuple[int, int] | None = None + + def __post_init__(self) -> None: + if not self.key: + raise ValueError("a row needs a store key") + if self.payload.overlaps(self.trailer): + raise ValueError(f"payload and trailer overlap for key {self.key}") + + +class RegionHandle(Protocol): + """Opaque handle of one registration; only the driver that issued it reads it.""" + + +class OpHandle(Protocol): + """Opaque handle of one put or get; only the driver that issued it reads it.""" + + +class ProbeStatus(enum.Enum): + """What a probe knows about one key, without moving payload.""" + + HIT = "hit" + MISS = "miss" + UNKNOWN = "unknown" + + +@dataclass(frozen=True) +class DriverCapabilities: + """What a driver implements beyond the required methods. + + Args: + cancel: ``cancel`` does something; without it the backend only suppresses unstarted work. + probe: ``probe`` answers; without it the backend fetches blindly. + max_rows_per_op: Rows one ``put``/``get`` accepts, or ``None`` for no limit. + """ + + cancel: bool = False + probe: bool = False + max_rows_per_op: int | None = None + + +@runtime_checkable +class StoreDriver(Protocol): + """What the backend asks of a store. + + Threads: ``register``, ``unregister``, ``put``, ``get``, ``cancel`` and ``drain`` are called one + at a time from the backend's own threads; ``poll`` and ``access_ended`` may be called from any + thread, concurrently with the rest. A driver over a blocking SDK blocks inside ``put``/``get``; + the backend keeps those calls off the submission thread. + + Side effects: ``register``, ``put`` and ``get`` raise ``ValueError`` only before anything + escaped, so the caller may forget the call. Once ``put``/``get`` has returned, every later + problem is a row result, and the operation holds the rows' memory until ``access_ended`` is + true. A failure return is not access-end evidence. + """ + + @property + def capabilities(self) -> DriverCapabilities: ... + + def register(self, address: int, nbytes: int, layout: RegionLayout) -> RegionHandle: + """Register one host region the driver may read and write. + + Raises: + ValueError: Bad geometry or overlap with a live registration; nothing escaped. + RegistrationError: The store refused; ``handle`` is set when the store may have + registered the region anyway. + """ + ... + + def unregister(self, handle: RegionHandle) -> None: + """Revoke a registration. Idempotent once it succeeded; a raise leaves the registration + live and the call retryable.""" + ... + + def put(self, rows: Sequence[RowIO]) -> OpHandle: + """Publish ``rows``; may return before any byte moved. + + Raises: + ValueError: A span lies outside a live registration, or ``rows`` exceed the driver's + capabilities; nothing escaped. + """ + ... + + def get(self, rows: Sequence[RowIO]) -> OpHandle: + """Fetch ``rows``; may return before any byte moved. Raises as ``put``.""" + ... + + def poll(self, op: OpHandle) -> tuple[RowResult, ...]: + """One result per row, in submission order; non-blocking. A concluded row never changes.""" + ... + + def access_ended(self, op: OpHandle) -> bool: + """True once the operation can no longer touch any of its rows' memory, with evidence. + Never turns false again. Independent of ``poll``.""" + ... + + def cancel(self, op: OpHandle) -> None: + """Best effort: rows not started conclude as CANCELLED; started rows may still finish. + Called only when ``capabilities.cancel`` is set.""" + ... + + def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...]: + """Advisory presence of ``keys``, within ``budget_s`` seconds; UNKNOWN when the budget runs + out. Called only when ``capabilities.probe`` is set. Raises on a store error rather than + answering MISS.""" + ... + + def drain(self, budget_s: float) -> bool: + """Stop starting work and wait up to ``budget_s`` seconds for every operation to end its + access. True when none can still touch memory; False means keep everything registered and + try again. Never forces an end.""" + ... diff --git a/tests/unittest/_torch/disaggregation/blob/host_memory.py b/tests/unittest/_torch/disaggregation/blob/host_memory.py new file mode 100644 index 000000000000..9bfb3fb2534c --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/host_memory.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host buffers with real addresses for driver tests, guarded by canaries on both sides.""" + +from __future__ import annotations + +import ctypes + +from tensorrt_llm._torch.disaggregation.blob.store import Span + +CANARY = bytes.fromhex("deadbeef") * 4 + + +class HostBuffer: + """``nbytes`` of host memory a driver may read and write, with a canary before and after. + + The buffer keeps its own reference to the backing storage, so the address stays valid for as + long as the object lives. + """ + + def __init__(self, nbytes: int, fill: int = 0) -> None: + if nbytes <= 0: + raise ValueError("a host buffer needs a positive size") + guard = len(CANARY) + self._raw = ctypes.create_string_buffer(guard * 2 + nbytes) + base = ctypes.addressof(self._raw) + ctypes.memmove(base, CANARY, guard) + ctypes.memmove(base + guard + nbytes, CANARY, guard) + self.address = base + guard + self.nbytes = nbytes + ctypes.memset(self.address, fill, nbytes) + + def span(self, offset: int = 0, nbytes: int | None = None) -> Span: + return Span(self.address + offset, self.nbytes - offset if nbytes is None else nbytes) + + def read(self, offset: int = 0, nbytes: int | None = None) -> bytes: + return ctypes.string_at( + self.address + offset, self.nbytes - offset if nbytes is None else nbytes + ) + + def write(self, data: bytes, offset: int = 0) -> None: + if offset < 0 or offset + len(data) > self.nbytes: + raise ValueError("write outside the buffer") + ctypes.memmove(self.address + offset, data, len(data)) + + def canaries_intact(self) -> bool: + guard = len(CANARY) + before = ctypes.string_at(self.address - guard, guard) + after = ctypes.string_at(self.address + self.nbytes, guard) + return before == CANARY and after == CANARY diff --git a/tests/unittest/_torch/disaggregation/blob/memory_store.py b/tests/unittest/_torch/disaggregation/blob/memory_store.py new file mode 100644 index 000000000000..eab2f9cf48c1 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/memory_store.py @@ -0,0 +1,373 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""An in-process store and a driver over it, for tests. + +``MemoryStore`` is a dictionary of complete objects that two endpoints in one process can share. +``MemoryDriver`` really copies bytes between registered host memory and that store. In automatic +mode an operation completes inside ``put``/``get``, like a blocking SDK. In controlled mode the test +releases three things separately per operation: the transfer, the logical results and the end of +access, so it can stage partial arrival, a result before or after access end, a late write after a +fault, and registration failures. No threads are started here. Transfers copy and call the store +outside the driver's lock, so queries never wait on the store; the lock is held around state +changes and around a late write's short copy, which keeps ``end_access`` out until it lands. +""" + +from __future__ import annotations + +import ctypes +import threading +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from typing import Literal + +from tensorrt_llm._torch.disaggregation.blob.format import TRAILER_BYTES +from tensorrt_llm._torch.disaggregation.blob.store import ( + CANCELLED, + MISS, + PENDING, + READY, + DriverCapabilities, + OpHandle, + ProbeStatus, + RegionHandle, + RegionLayout, + RegistrationError, + RowIO, + RowResult, + RowStatus, + Span, + StoreError, + faulted, +) + + +class MemoryStore: + """Objects by key. Each object is a private copy, committed whole or not at all. + + Subclass and override ``commit`` or ``get`` to make a test's store block or fail; both may + raise ``StoreError``, which the driver turns into a row fault. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._objects: dict[str, bytes] = {} + + def commit(self, key: str, data: bytes) -> bool: + """Store ``data`` under ``key``. False when a different object already holds the key; + the existing object is kept. Storing identical bytes again is a no-op success.""" + copy = bytes(data) + with self._lock: + existing = self._objects.get(key) + if existing is None: + self._objects[key] = copy + return True + return existing == copy + + def get(self, key: str) -> bytes | None: + with self._lock: + return self._objects.get(key) + + def contains(self, key: str) -> bool: + return self.get(key) is not None + + def committed_object(self, key: str) -> bytes: + """The whole object under ``key``; ``KeyError`` when nothing is committed there.""" + data = self.get(key) + if data is None: + raise KeyError(f"no committed object under {key!r}") + return data + + def committed_payload(self, key: str, trailer_nbytes: int = TRAILER_BYTES) -> bytes: + """The payload of the object under ``key``, without its trailer.""" + return self.committed_object(key)[:-trailer_nbytes] + + def __len__(self) -> int: + with self._lock: + return len(self._objects) + + +@dataclass(eq=False) +class _Registration: + span: Span + layout: RegionLayout + live: bool = True + + +@dataclass(eq=False) +class _Row: + io: RowIO + published: RowResult = PENDING + staged: RowResult | None = None # decided but not yet visible through poll + started: bool = False # a transfer took the row; cancel and fail leave it alone + transferred: bool = False + + +@dataclass(eq=False) +class MemoryOp: + """One put or get. Tests hold it to drive the controlled mode.""" + + kind: Literal["put", "get"] + rows: list[_Row] + access_ended: bool = False + + +class MemoryDriver: + """A ``StoreDriver`` over a ``MemoryStore``; see the module docstring for the two modes.""" + + def __init__( + self, + store: MemoryStore, + *, + controlled: bool = False, + capabilities: DriverCapabilities = DriverCapabilities(cancel=True, probe=True), + ) -> None: + self.store = store + self.controlled = controlled + self._capabilities = capabilities + self._lock = threading.Condition(threading.RLock()) + self._registrations: list[_Registration] = [] + self._ops: list[MemoryOp] = [] + self._draining = False + self._cancel_effective = True + self._probe_unknown = False + self._register_failure: tuple[str, bool] | None = None + self._unregister_failures: dict[int, str] = {} + + # The StoreDriver protocol. + + @property + def capabilities(self) -> DriverCapabilities: + return self._capabilities + + def register(self, address: int, nbytes: int, layout: RegionLayout) -> _Registration: + span = Span(address, nbytes) + if layout.nbytes > nbytes: + raise ValueError(f"layout needs {layout.nbytes} bytes, region has {nbytes}") + with self._lock: + for other in self._registrations: + if other.live and other.span.overlaps(span): + raise ValueError(f"region {span} overlaps a live registration {other.span}") + registration = _Registration(span, layout) + if self._register_failure is not None: + reason, escaped = self._register_failure + self._register_failure = None + if escaped: + # The store may have registered before it answered: hand out the handle. + self._registrations.append(registration) + raise RegistrationError(reason, registration) + raise RegistrationError(reason) + self._registrations.append(registration) + return registration + + def unregister(self, handle: RegionHandle) -> None: + with self._lock: + registration = self._registration(handle) + reason = self._unregister_failures.pop(id(registration), None) + if reason is not None: + raise StoreError(reason) + registration.live = False + + def put(self, rows: Sequence[RowIO]) -> MemoryOp: + return self._submit("put", rows) + + def get(self, rows: Sequence[RowIO]) -> MemoryOp: + return self._submit("get", rows) + + def poll(self, op: OpHandle) -> tuple[RowResult, ...]: + with self._lock: + return tuple(row.published for row in self._op(op).rows) + + def access_ended(self, op: OpHandle) -> bool: + with self._lock: + return self._op(op).access_ended + + def cancel(self, op: OpHandle) -> None: + with self._lock: + if self._cancel_effective: + self._cancel_unstarted(self._op(op)) + + def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...]: + if self._probe_unknown: + return tuple(ProbeStatus.UNKNOWN for _ in keys) + return tuple( + ProbeStatus.HIT if self.store.contains(key) else ProbeStatus.MISS for key in keys + ) + + def drain(self, budget_s: float) -> bool: + with self._lock: + self._draining = True + for op in self._ops: + self._cancel_unstarted(op) + return self._lock.wait_for( + lambda: all(op.access_ended for op in self._ops), timeout=budget_s + ) + + # Controlled mode: the test releases each stage. + + def transfer(self, op: MemoryOp, rows: Iterable[int] | None = None) -> None: + """Move the bytes of ``rows`` (all by default) and stage their results. + + The copies and store calls run outside the driver's lock; a ``StoreError`` from the store + becomes that row's fault instead of escaping. + """ + with self._lock: + self._require_open(op) + taken = [] + for index in self._indexes(op, rows): + row = op.rows[index] + if row.started or row.published.concluded: + raise RuntimeError(f"row {index} was already taken or decided") + row.started = True + taken.append(row) + for row in taken: + result = self._move(op.kind, row.io) + with self._lock: + row.staged = result + row.transferred = True + + def publish_results(self, op: MemoryOp, rows: Iterable[int] | None = None) -> None: + """Make the staged results of ``rows`` visible through ``poll``.""" + with self._lock: + for index in self._indexes(op, rows): + row = op.rows[index] + if row.staged is None: + raise RuntimeError(f"row {index} has no staged result to publish") + row.published, row.staged = row.staged, None + + def complete(self, op: MemoryOp, rows: Iterable[int] | None = None) -> None: + """Transfer ``rows`` and publish their results, leaving access open.""" + indexes = self._indexes(op, rows) # ``rows`` may be a one-shot iterator + self.transfer(op, indexes) + self.publish_results(op, indexes) + + def fail(self, op: MemoryOp, reason: str, rows: Iterable[int] | None = None) -> None: + """Conclude ``rows`` as FAULT without moving bytes; access stays open.""" + with self._lock: + self._require_open(op) + for index in self._indexes(op, rows): + row = op.rows[index] + if row.started or row.published.concluded: + raise RuntimeError(f"row {index} was already taken or decided") + row.published = faulted(reason) + + def write_late(self, op: MemoryOp, index: int, data: bytes) -> None: + """Write into a faulted get row's payload after its result, as a transfer that kept going + would. Refused once access has ended, and for rows that are READY, MISS, CANCELLED or + still pending: those must stay as reported.""" + with self._lock: + self._require_open(op) + if op.kind != "get": + raise RuntimeError("late writes only happen on fetch destinations") + row = op.rows[index] + if row.published.status is not RowStatus.FAULT: + raise RuntimeError(f"row {index} is {row.published.status.value}, not faulted") + payload = row.io.payload + if len(data) != payload.nbytes: + raise ValueError( + f"late write of {len(data)} bytes into a {payload.nbytes}-byte row" + ) + ctypes.memmove(payload.address, data, len(data)) + + def end_access(self, op: MemoryOp) -> None: + """Declare that the operation will never touch its rows' memory again.""" + with self._lock: + if any(row.started and not row.transferred for row in op.rows): + raise RuntimeError("a transfer is still in progress") + op.access_ended = True + self._lock.notify_all() + + def set_cancel_effective(self, effective: bool) -> None: + self._cancel_effective = effective + + def set_probe_unknown(self, unknown: bool) -> None: + self._probe_unknown = unknown + + def fail_next_register(self, reason: str, *, escaped: bool) -> None: + """Make the next ``register`` raise; with ``escaped`` the store keeps the region anyway.""" + self._register_failure = (reason, escaped) + + def fail_next_unregister(self, handle: RegionHandle, reason: str) -> None: + with self._lock: + self._unregister_failures[id(self._registration(handle))] = reason + + def live_registrations(self) -> list[_Registration]: + with self._lock: + return [r for r in self._registrations if r.live] + + def open_ops(self) -> list[MemoryOp]: + with self._lock: + return [op for op in self._ops if not op.access_ended] + + # Internals. + + def _submit(self, kind: Literal["put", "get"], rows: Sequence[RowIO]) -> MemoryOp: + if not rows: + raise ValueError("an operation needs at least one row") + limit = self._capabilities.max_rows_per_op + if limit is not None and len(rows) > limit: + raise ValueError(f"{len(rows)} rows exceed the limit of {limit} per operation") + with self._lock: + if self._draining: + raise ValueError("the driver is draining") + for io in rows: + for span in (io.payload, io.trailer): + if not self._registered(span): + raise ValueError(f"{span} is outside every live registration") + op = MemoryOp(kind, [_Row(io) for io in rows]) + self._ops.append(op) + if not self.controlled: + # Like a blocking SDK: the call returns once every row is decided and access is over. + self.transfer(op) + self.publish_results(op) + self.end_access(op) + return op + + def _move(self, kind: str, io: RowIO) -> RowResult: + try: + return self._put_row(io) if kind == "put" else self._get_row(io) + except StoreError as error: + return faulted(str(error)) + + def _put_row(self, io: RowIO) -> RowResult: + data = ctypes.string_at(io.payload.address, io.payload.nbytes) + data += ctypes.string_at(io.trailer.address, io.trailer.nbytes) + if self.store.commit(io.key, data): + return READY + return faulted("key exists with different content") + + def _get_row(self, io: RowIO) -> RowResult: + data = self.store.get(io.key) + if data is None: + return MISS + expected = io.payload.nbytes + io.trailer.nbytes + if len(data) != expected: + return faulted(f"object is {len(data)} bytes, expected {expected}") + ctypes.memmove(io.payload.address, data, io.payload.nbytes) + ctypes.memmove(io.trailer.address, data[io.payload.nbytes :], io.trailer.nbytes) + return READY + + def _cancel_unstarted(self, op: MemoryOp) -> None: + for row in op.rows: + if not row.started and not row.published.concluded: + row.published = CANCELLED + + def _registered(self, span: Span) -> bool: + return any(r.live and span.within(r.span) for r in self._registrations) + + def _registration(self, handle: RegionHandle) -> _Registration: + if not isinstance(handle, _Registration) or handle not in self._registrations: + raise ValueError("handle does not belong to this driver") + return handle + + def _op(self, op: OpHandle) -> MemoryOp: + if not isinstance(op, MemoryOp) or op not in self._ops: + raise ValueError("operation does not belong to this driver") + return op + + def _require_open(self, op: MemoryOp) -> None: + if self._op(op).access_ended: + raise RuntimeError("the operation already ended its access") + + @staticmethod + def _indexes(op: MemoryOp, rows: Iterable[int] | None) -> list[int]: + return list(range(len(op.rows))) if rows is None else list(rows) diff --git a/tests/unittest/_torch/disaggregation/blob/test_memory_store.py b/tests/unittest/_torch/disaggregation/blob/test_memory_store.py new file mode 100644 index 000000000000..237024e40751 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_memory_store.py @@ -0,0 +1,566 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The private driver boundary, exercised through the in-memory store and driver. + +These tests pin what any driver owes the backend: bytes move whole and never alias the source, +results are per row and never pass a timeout off as a miss, access end is separate evidence, and +failed registrations stay cleanable. +""" + +from __future__ import annotations + +import threading + +import pytest +from host_memory import HostBuffer +from memory_store import MemoryDriver, MemoryOp, MemoryStore + +from tensorrt_llm._torch.disaggregation.blob.format import ( + TRAILER_BYTES, + check_trailer, + encode_trailer, +) +from tensorrt_llm._torch.disaggregation.blob.keys import KEY_FORMAT_VERSION, encode_key +from tensorrt_llm._torch.disaggregation.blob.store import ( + ProbeStatus, + RegionLayout, + RegistrationError, + RowIO, + RowStatus, + Span, + StoreDriver, + StoreError, + concluded, +) + +pytestmark = pytest.mark.cpu_only + +SLOT = 8 +NAMES = {"A": b"name-A", "B": b"name-B", "C": b"name-C"} + + +class Endpoint: + """One side of a transfer: a registered payload region plus registered trailer buffers.""" + + def __init__(self, driver: MemoryDriver, labels: list[str], fill: int = 0) -> None: + self.driver = driver + self.labels = labels + self.payload = HostBuffer(SLOT * len(labels), fill) + self.trailer = HostBuffer(TRAILER_BYTES * len(labels), fill) + self.handles = [ + driver.register( + self.payload.address, self.payload.nbytes, RegionLayout(SLOT, len(labels), "kv") + ), + driver.register( + self.trailer.address, + self.trailer.nbytes, + RegionLayout(TRAILER_BYTES, len(labels), "trailer"), + ), + ] + self.rows = [ + RowIO( + encode_key("t", NAMES[label]), + self.payload.span(i * SLOT, SLOT), + self.trailer.span(i * TRAILER_BYTES, TRAILER_BYTES), + (0, i), + ) + for i, label in enumerate(labels) + ] + + def fill_row(self, index: int, data: bytes) -> None: + self.payload.write(data, index * SLOT) + self.trailer.write(encode_trailer(NAMES[self.labels[index]], SLOT), index * TRAILER_BYTES) + + def read_row(self, index: int) -> bytes: + return self.payload.read(index * SLOT, SLOT) + + def trailer_ok(self, index: int) -> str | None: + trailer = self.trailer.read(index * TRAILER_BYTES, TRAILER_BYTES) + return check_trailer(trailer, NAMES[self.labels[index]], SLOT) + + def intact(self) -> bool: + return self.payload.canaries_intact() and self.trailer.canaries_intact() + + +def publish(store: MemoryStore, labels: dict[str, bytes]) -> None: + """Put ``labels`` into ``store`` through an automatic driver of its own.""" + writer = Endpoint(MemoryDriver(store), list(labels)) + for i, data in enumerate(labels.values()): + writer.fill_row(i, data) + op = writer.driver.put(writer.rows) + assert all(r.status is RowStatus.READY for r in writer.driver.poll(op)) + + +@pytest.fixture +def store() -> MemoryStore: + return MemoryStore() + + +def statuses(driver: MemoryDriver, op: MemoryOp) -> list[RowStatus]: + return [r.status for r in driver.poll(op)] + + +# Bytes and ownership. + + +def test_put_keeps_its_own_copy_so_the_source_may_change(store: MemoryStore) -> None: + writer = Endpoint(MemoryDriver(store), ["A"]) + writer.fill_row(0, b"aaaaaaaa") + op = writer.driver.put(writer.rows) + assert statuses(writer.driver, op) == [RowStatus.READY] + assert writer.driver.access_ended(op) + + writer.payload.write(b"xxxxxxxx") + assert store.committed_payload(writer.rows[0].key) == b"aaaaaaaa" + + reader = Endpoint(MemoryDriver(store), ["A"], fill=0x2D) + got = reader.driver.get(reader.rows) + assert statuses(reader.driver, got) == [RowStatus.READY] + assert reader.read_row(0) == b"aaaaaaaa" + assert reader.trailer_ok(0) is None + assert reader.intact() + + +def test_get_reports_hit_hit_miss_per_row_and_leaves_the_miss_untouched(store: MemoryStore) -> None: + publish(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + reader = Endpoint(MemoryDriver(store), ["A", "B", "C"], fill=0x2D) + op = reader.driver.get(reader.rows) + assert statuses(reader.driver, op) == [RowStatus.READY, RowStatus.READY, RowStatus.MISS] + assert concluded(reader.driver.poll(op)) + assert reader.read_row(0) == b"aaaaaaaa" + assert reader.read_row(1) == b"bbbbbbbb" + assert reader.read_row(2) == b"-" * SLOT + assert reader.trailer_ok(2) is not None + assert reader.intact() + + +def test_get_rejects_a_size_mismatch_before_writing(store: MemoryStore) -> None: + key = encode_key("t", NAMES["A"]) + assert store.commit(key, b"short") + reader = Endpoint(MemoryDriver(store), ["A"], fill=0x2D) + op = reader.driver.get(reader.rows) + (result,) = reader.driver.poll(op) + assert result.status is RowStatus.FAULT + assert result.fault == f"object is 5 bytes, expected {SLOT + TRAILER_BYTES}" + assert reader.read_row(0) == b"-" * SLOT + assert reader.intact() + + +def test_rows_outside_a_live_registration_are_refused_before_any_effect(store: MemoryStore) -> None: + driver = MemoryDriver(store) + ok = Endpoint(driver, ["A"]) + stray = HostBuffer(SLOT) + row = RowIO(ok.rows[0].key, stray.span(), ok.rows[0].trailer) + with pytest.raises(ValueError, match="outside"): + driver.get([row]) + assert driver.open_ops() == [] + driver.unregister(ok.handles[0]) + with pytest.raises(ValueError, match="outside"): + driver.get(ok.rows) + + +# Format. + + +def test_trailer_round_trip_and_mismatches() -> None: + trailer = encode_trailer(b"name", 64) + assert len(trailer) == TRAILER_BYTES + assert check_trailer(trailer, b"name", 64) is None + assert check_trailer(trailer, b"other", 64) == "object name mismatch" + assert check_trailer(trailer, b"name", 65) == "payload length 64, expected 65" + assert check_trailer(b"\0" * TRAILER_BYTES, b"name", 64) == "bad magic" + assert check_trailer(trailer[:-1], b"name", 64) == ( + f"trailer is {TRAILER_BYTES - 1} bytes, expected {TRAILER_BYTES}" + ) + + +def test_key_encoding_is_prefix_version_and_hex() -> None: + key = encode_key("prod", b"\x01\xff") + assert key == f"prod/v{KEY_FORMAT_VERSION}/01ff" + for prefix in ("", "a/b", "a b"): + with pytest.raises(ValueError): + encode_key(prefix, b"\x01") + with pytest.raises(ValueError): + encode_key("prod", b"") + + +def test_row_refuses_overlapping_payload_and_trailer() -> None: + with pytest.raises(ValueError, match="overlap"): + RowIO("k", Span(0x1000, 64), Span(0x1020, 32)) + + +# Publication. + + +def test_duplicate_put_reuses_identical_object_and_rejects_different_bytes( + store: MemoryStore, +) -> None: + first = Endpoint(MemoryDriver(store), ["A"]) + first.fill_row(0, b"aaaaaaaa") + assert statuses(first.driver, first.driver.put(first.rows)) == [RowStatus.READY] + + same = Endpoint(MemoryDriver(store), ["A"]) + same.fill_row(0, b"aaaaaaaa") + assert statuses(same.driver, same.driver.put(same.rows)) == [RowStatus.READY] + + other = Endpoint(MemoryDriver(store), ["A"]) + other.fill_row(0, b"zzzzzzzz") + (result,) = other.driver.poll(other.driver.put(other.rows)) + assert result.status is RowStatus.FAULT + assert store.committed_payload(first.rows[0].key) == b"aaaaaaaa" + assert len(store) == 1 + + +def test_a_failed_retry_does_not_damage_the_committed_object(store: MemoryStore) -> None: + publish(store, {"A": b"aaaaaaaa"}) + driver = MemoryDriver(store, controlled=True) + retry = Endpoint(driver, ["A"]) + retry.fill_row(0, b"aaaaaaaa") + op = driver.put(retry.rows) + driver.fail(op, "connection reset") + (result,) = driver.poll(op) + assert result.status is RowStatus.FAULT + assert not driver.access_ended(op) + assert store.committed_payload(retry.rows[0].key) == b"aaaaaaaa" + + +def test_readers_never_see_a_partial_object(store: MemoryStore) -> None: + pattern_a, pattern_b = b"a" * 4096, b"b" * 4096 + keys = [f"k{i}" for i in range(200)] + stop = threading.Event() + seen_partial = [] + + def write() -> None: + for i, key in enumerate(keys): + store.commit(key, pattern_a if i % 2 == 0 else pattern_b) + stop.set() + + def read() -> None: + while not stop.is_set(): + for key in keys: + data = store.get(key) + if data is not None and data not in (pattern_a, pattern_b): + seen_partial.append(key) + + readers = [threading.Thread(target=read) for _ in range(2)] + writer = threading.Thread(target=write) + for t in readers + [writer]: + t.start() + for t in readers + [writer]: + t.join() + assert seen_partial == [] + assert len(store) == len(keys) + + +# Result and access end are independent. + + +def test_result_can_come_before_access_end(store: MemoryStore) -> None: + publish(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A", "B"], fill=0x2D) + op = driver.get(reader.rows) + assert statuses(driver, op) == [RowStatus.PENDING, RowStatus.PENDING] + + driver.complete(op, [0]) + assert statuses(driver, op) == [RowStatus.READY, RowStatus.PENDING] + assert reader.read_row(0) == b"aaaaaaaa" + assert reader.read_row(1) == b"-" * SLOT + + driver.fail(op, "timeout", [1]) + results = driver.poll(op) + assert concluded(results) + assert results[1].status is RowStatus.FAULT and results[1].fault == "timeout" + assert not driver.access_ended(op) + + driver.end_access(op) + assert driver.access_ended(op) + assert driver.poll(op) == results + + +def test_complete_accepts_a_one_shot_iterator(store: MemoryStore) -> None: + publish(store, {"A": b"aaaaaaaa"}) + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A"], fill=0x2D) + op = driver.get(reader.rows) + driver.complete(op, iter([0])) + assert statuses(driver, op) == [RowStatus.READY] + assert reader.read_row(0) == b"aaaaaaaa" + + +def test_access_can_end_before_the_result_is_published(store: MemoryStore) -> None: + publish(store, {"A": b"aaaaaaaa"}) + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A"], fill=0x2D) + op = driver.get(reader.rows) + driver.transfer(op) + driver.end_access(op) + assert driver.access_ended(op) + assert statuses(driver, op) == [RowStatus.PENDING] + driver.publish_results(op) + assert statuses(driver, op) == [RowStatus.READY] + assert reader.read_row(0) == b"aaaaaaaa" + + +def test_a_late_write_lands_only_while_access_is_open(store: MemoryStore) -> None: + publish(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A", "B"], fill=0x2D) + op = driver.get(reader.rows) + driver.complete(op, [0]) + driver.fail(op, "timeout", [1]) + results = driver.poll(op) + + driver.write_late(op, 1, b"bbbbbbbb") + assert reader.read_row(1) == b"bbbbbbbb" + assert driver.poll(op) == results + assert not driver.access_ended(op) + + with pytest.raises(RuntimeError, match="ready, not faulted"): + driver.write_late(op, 0, b"cccccccc") + assert reader.read_row(0) == b"aaaaaaaa" + + driver.end_access(op) + with pytest.raises(RuntimeError, match="ended"): + driver.write_late(op, 1, b"cccccccc") + with pytest.raises(RuntimeError, match="ended"): + driver.transfer(op, [1]) + assert reader.read_row(1) == b"bbbbbbbb" + assert reader.intact() + + +def test_a_cancelled_row_cannot_be_written_late(store: MemoryStore) -> None: + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A"], fill=0x2D) + op = driver.get(reader.rows) + driver.cancel(op) + assert statuses(driver, op) == [RowStatus.CANCELLED] + with pytest.raises(RuntimeError, match="cancelled, not faulted"): + driver.write_late(op, 0, b"aaaaaaaa") + assert reader.read_row(0) == b"-" * SLOT + + +def test_timeout_is_a_fault_not_a_miss(store: MemoryStore) -> None: + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["C"]) + op = driver.get(reader.rows) + driver.fail(op, "timeout") + (result,) = driver.poll(op) + assert result.status is RowStatus.FAULT + assert result.fault == "timeout" + + +class FailingStore(MemoryStore): + """Raises ``StoreError`` on the n-th commit, as a store that drops a connection would.""" + + def __init__(self, fail_on_commit: int) -> None: + super().__init__() + self._fail_on = fail_on_commit + self._commits = 0 + + def commit(self, key: str, data: bytes) -> bool: + self._commits += 1 + if self._commits == self._fail_on: + raise StoreError("connection dropped") + return super().commit(key, data) + + +def test_a_store_error_mid_operation_becomes_a_row_fault_and_keeps_the_handle() -> None: + store = FailingStore(fail_on_commit=2) + writer = Endpoint(MemoryDriver(store), ["A", "B", "C"]) + for i, data in enumerate((b"aaaaaaaa", b"bbbbbbbb", b"cccccccc")): + writer.fill_row(i, data) + op = writer.driver.put(writer.rows) + results = writer.driver.poll(op) + assert [r.status for r in results] == [RowStatus.READY, RowStatus.FAULT, RowStatus.READY] + assert results[1].fault == "connection dropped" + assert writer.driver.access_ended(op) + assert store.committed_payload(writer.rows[0].key) == b"aaaaaaaa" + assert not store.contains(writer.rows[1].key) + assert store.committed_payload(writer.rows[2].key) == b"cccccccc" + assert writer.driver.drain(0.0) is True + + +class GatedStore(MemoryStore): + """Holds every commit until the test opens the gate, and says when one is waiting.""" + + def __init__(self) -> None: + super().__init__() + self.gate = threading.Event() + self.waiting = threading.Event() + + def commit(self, key: str, data: bytes) -> bool: + self.waiting.set() + self.gate.wait() + return super().commit(key, data) + + +def test_queries_do_not_wait_on_another_operation_in_the_store() -> None: + store = GatedStore() + driver = MemoryDriver(store, controlled=True) + done = Endpoint(driver, ["A"]) + done.fill_row(0, b"aaaaaaaa") + finished = driver.get(done.rows) + driver.fail(finished, "timeout") + + slow = Endpoint(driver, ["B"]) + slow.fill_row(0, b"bbbbbbbb") + blocked = driver.put(slow.rows) + worker = threading.Thread(target=driver.transfer, args=(blocked,), daemon=True) + worker.start() + answered = threading.Event() + seen: dict[str, object] = {} + failure: list[Exception] = [] + + def query() -> None: + # Runs off the main thread: if a query waits on the store, the main thread can still + # open the gate and the test fails instead of hanging. + try: + seen["finished"] = statuses(driver, finished) + seen["blocked"] = statuses(driver, blocked) + seen["access_ended"] = driver.access_ended(blocked) + try: + driver.end_access(blocked) + except RuntimeError as error: + seen["end_access"] = str(error) + driver.cancel(blocked) + seen["after_cancel"] = statuses(driver, blocked) + except Exception as error: # surfaced on the main thread below + failure.append(error) + finally: + answered.set() + + querier = threading.Thread(target=query, daemon=True) + try: + assert store.waiting.wait(5.0), "the worker never reached the store" + # The worker is inside the store with its row taken; these must answer at once. + querier.start() + assert answered.wait(5.0), "queries waited on an operation that is inside the store" + finally: + store.gate.set() + worker.join(5.0) + if querier.ident is not None: # started; joining an unstarted thread raises + querier.join(5.0) + assert not worker.is_alive() and not querier.is_alive() + if failure: + raise AssertionError("a query raised") from failure[0] + assert seen["finished"] == [RowStatus.FAULT] + assert seen["blocked"] == [RowStatus.PENDING] + assert seen["access_ended"] is False + assert seen.get("end_access") == "a transfer is still in progress" + assert seen["after_cancel"] == [RowStatus.PENDING] + driver.publish_results(blocked) + assert statuses(driver, blocked) == [RowStatus.READY] + assert store.committed_payload(slow.rows[0].key) == b"bbbbbbbb" + + +# Cancellation, probe, drain. + + +def test_cancel_concludes_unstarted_rows_and_spares_started_ones(store: MemoryStore) -> None: + publish(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A", "B"], fill=0x2D) + op = driver.get(reader.rows) + driver.transfer(op, [0]) + driver.cancel(op) + assert statuses(driver, op) == [RowStatus.PENDING, RowStatus.CANCELLED] + driver.publish_results(op, [0]) + assert statuses(driver, op) == [RowStatus.READY, RowStatus.CANCELLED] + assert reader.read_row(1) == b"-" * SLOT + assert not driver.access_ended(op) + + +def test_an_ineffective_cancel_changes_nothing(store: MemoryStore) -> None: + driver = MemoryDriver(store, controlled=True) + driver.set_cancel_effective(False) + reader = Endpoint(driver, ["A"]) + op = driver.get(reader.rows) + driver.cancel(op) + assert statuses(driver, op) == [RowStatus.PENDING] + + +def test_probe_answers_hit_miss_or_unknown(store: MemoryStore) -> None: + publish(store, {"A": b"aaaaaaaa"}) + driver = MemoryDriver(store) + keys = [encode_key("t", NAMES["A"]), encode_key("t", NAMES["C"])] + assert driver.probe(keys, 1.0) == (ProbeStatus.HIT, ProbeStatus.MISS) + driver.set_probe_unknown(True) + assert driver.probe(keys, 1.0) == (ProbeStatus.UNKNOWN, ProbeStatus.UNKNOWN) + + +def test_drain_waits_for_access_end_and_refuses_new_work(store: MemoryStore) -> None: + publish(store, {"A": b"aaaaaaaa"}) + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A"]) + op = driver.get(reader.rows) + assert driver.drain(0.01) is False + assert statuses(driver, op) == [RowStatus.CANCELLED] + with pytest.raises(ValueError, match="draining"): + driver.get(reader.rows) + + release = threading.Timer(0.05, driver.end_access, args=(op,)) + release.start() + try: + assert driver.drain(5.0) is True + finally: + release.join() + assert driver.open_ops() == [] + + +# Registration. + + +def test_register_checks_geometry_and_overlap_and_unregister_is_idempotent( + store: MemoryStore, +) -> None: + driver = MemoryDriver(store) + buffer = HostBuffer(64) + with pytest.raises(ValueError, match="layout needs"): + driver.register(buffer.address, 64, RegionLayout(SLOT, 9)) + handle = driver.register(buffer.address, 64, RegionLayout(SLOT, 8)) + with pytest.raises(ValueError, match="overlaps"): + driver.register(buffer.address + 8, 8, RegionLayout(8, 1)) + driver.unregister(handle) + driver.unregister(handle) + assert driver.live_registrations() == [] + driver.register(buffer.address + 8, 8, RegionLayout(8, 1)) + + +def test_a_registration_the_store_may_hold_keeps_its_handle(store: MemoryStore) -> None: + driver = MemoryDriver(store) + buffer = HostBuffer(64) + driver.fail_next_register("store unreachable", escaped=False) + with pytest.raises(RegistrationError) as clean: + driver.register(buffer.address, 64, RegionLayout(SLOT, 8)) + assert clean.value.handle is None + assert driver.live_registrations() == [] + + driver.fail_next_register("timed out after the request", escaped=True) + with pytest.raises(RegistrationError) as ambiguous: + driver.register(buffer.address, 64, RegionLayout(SLOT, 8)) + handle = ambiguous.value.handle + assert handle is not None + assert driver.live_registrations() == [handle] + + driver.fail_next_unregister(handle, "busy") + with pytest.raises(StoreError, match="busy"): + driver.unregister(handle) + assert driver.live_registrations() == [handle] + driver.unregister(handle) + assert driver.live_registrations() == [] + + +def test_memory_driver_has_every_member_of_the_driver_protocol(store: MemoryStore) -> None: + driver: StoreDriver = MemoryDriver(store) + assert isinstance(driver, StoreDriver) + + +def test_automatic_mode_follows_the_same_protocol(store: MemoryStore) -> None: + writer = Endpoint(MemoryDriver(store), ["A"]) + writer.fill_row(0, b"aaaaaaaa") + op = writer.driver.put(writer.rows) + assert concluded(writer.driver.poll(op)) + assert writer.driver.access_ended(op) + assert writer.driver.drain(0.0) is True + assert writer.driver.open_ops() == [] From e126bc9b9b945f78a472943bddc661b35e58241e Mon Sep 17 00:00:00 2001 From: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> Date: Thu, 15 Oct 2026 18:09:33 +0800 Subject: [PATCH 2/6] [None][feat] Leave object format to drivers and prove each driver check bites Rows carry key and payload only; drivers own format and integrity, and NOT_READY and DECLINED are results rather than faults. Driver checks run through a factory with cleanup, and each is proven by a driver that breaks exactly its rule. Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> --- .../_torch/disaggregation/blob/__init__.py | 5 +- .../_torch/disaggregation/blob/format.py | 49 - .../_torch/disaggregation/blob/keys.py | 23 +- .../_torch/disaggregation/blob/store.py | 163 ++- .../disaggregation/blob/driver_conformance.py | 503 ++++++++ .../disaggregation/blob/memory_store.py | 405 ++++-- .../disaggregation/blob/object_trailer.py | 98 ++ .../disaggregation/blob/test_blob_format.py | 135 ++ .../disaggregation/blob/test_memory_driver.py | 1081 +++++++++++++++++ .../disaggregation/blob/test_memory_store.py | 566 --------- 10 files changed, 2289 insertions(+), 739 deletions(-) delete mode 100644 tensorrt_llm/_torch/disaggregation/blob/format.py create mode 100644 tests/unittest/_torch/disaggregation/blob/driver_conformance.py create mode 100644 tests/unittest/_torch/disaggregation/blob/object_trailer.py create mode 100644 tests/unittest/_torch/disaggregation/blob/test_blob_format.py create mode 100644 tests/unittest/_torch/disaggregation/blob/test_memory_driver.py delete mode 100644 tests/unittest/_torch/disaggregation/blob/test_memory_store.py diff --git a/tensorrt_llm/_torch/disaggregation/blob/__init__.py b/tensorrt_llm/_torch/disaggregation/blob/__init__.py index 7c7756c302ba..5e7ae28d4d83 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/blob/__init__.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """The blob backend: KV rows shared through a byte store. -Only the private driver boundary lives here so far. Nothing in this package imports the public -backend contract or an optional store SDK at import time. +So far this package holds the private driver boundary (``store.py``) and the key encoding +(``keys.py``). Nothing here imports the public backend contract or an optional store SDK at +import time. """ diff --git a/tensorrt_llm/_torch/disaggregation/blob/format.py b/tensorrt_llm/_torch/disaggregation/blob/format.py deleted file mode 100644 index 9875fb8227f8..000000000000 --- a/tensorrt_llm/_torch/disaggregation/blob/format.py +++ /dev/null @@ -1,49 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The object format: the payload bytes followed by a fixed-size trailer. - -Private and versioned. The trailer lets a reader check that an object has the expected format, -belongs to the expected name and has the expected length; it is not a payload checksum and proves -neither a complete commit nor the end of access. Any incompatible change bumps ``FORMAT_VERSION``. -""" - -from __future__ import annotations - -import hashlib -import struct - -FORMAT_VERSION = 1 -MAGIC = b"TLBK" - -# magic, version, 2 bytes of padding, name digest, payload length. -_TRAILER = struct.Struct("<4sH2x16sQ") -TRAILER_BYTES = _TRAILER.size -_DIGEST_BYTES = 16 - - -def name_digest(name: bytes) -> bytes: - """The 16-byte digest of an object name that the trailer carries.""" - return hashlib.blake2b(name, digest_size=_DIGEST_BYTES).digest() - - -def encode_trailer(name: bytes, payload_nbytes: int) -> bytes: - """The trailer for an object of ``payload_nbytes`` bytes stored under ``name``.""" - if payload_nbytes <= 0: - raise ValueError(f"payload length must be positive, got {payload_nbytes}") - return _TRAILER.pack(MAGIC, FORMAT_VERSION, name_digest(name), payload_nbytes) - - -def check_trailer(trailer: bytes, name: bytes, payload_nbytes: int) -> str | None: - """``None`` when ``trailer`` matches ``name`` and ``payload_nbytes``, else why it does not.""" - if len(trailer) != TRAILER_BYTES: - return f"trailer is {len(trailer)} bytes, expected {TRAILER_BYTES}" - magic, version, digest, length = _TRAILER.unpack(trailer) - if magic != MAGIC: - return "bad magic" - if version != FORMAT_VERSION: - return f"format version {version}, expected {FORMAT_VERSION}" - if digest != name_digest(name): - return "object name mismatch" - if length != payload_nbytes: - return f"payload length {length}, expected {payload_nbytes}" - return None diff --git a/tensorrt_llm/_torch/disaggregation/blob/keys.py b/tensorrt_llm/_torch/disaggregation/blob/keys.py index c8fbc39f8722..a27af1c7d481 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/keys.py +++ b/tensorrt_llm/_torch/disaggregation/blob/keys.py @@ -2,27 +2,32 @@ # SPDX-License-Identifier: Apache-2.0 """Store keys for blob objects. -A private, versioned encoding of the lender's opaque object names, not a stable public format. An -incompatible change bumps ``KEY_FORMAT_VERSION`` so objects written under the old one miss instead -of being misread. +A private encoding of the lender's opaque object names, not a stable public format. The key +carries ``FORMAT_VERSION``, which covers the key spelling and the payload layout alike: after an +incompatible change, objects written under the old version are never looked up, so they miss +instead of being misread. """ from __future__ import annotations -KEY_FORMAT_VERSION = 1 +import re +FORMAT_VERSION = 1 + +_PREFIX = re.compile(r"[A-Za-z0-9_-]{1,64}") _SEPARATOR = "/" def encode_key(prefix: str, name: bytes) -> str: - """``"/v/"``. + """``"/v/"``. Args: - prefix: The deployment's namespace; non-empty, without ``/`` or whitespace. + prefix: The store partition the objects live in, for example a deployment name; matches + ``[A-Za-z0-9_-]{1,64}``. name: The lender's opaque object name; encoded whole, never parsed. """ - if not prefix or _SEPARATOR in prefix or any(c.isspace() for c in prefix): - raise ValueError(f"bad key prefix {prefix!r}: non-empty, no '/' and no whitespace") + if not _PREFIX.fullmatch(prefix): + raise ValueError(f"bad key prefix {prefix!r}: expected [A-Za-z0-9_-]{{1,64}}") if not name: raise ValueError("an object name is non-empty") - return f"{prefix}{_SEPARATOR}v{KEY_FORMAT_VERSION}{_SEPARATOR}{name.hex()}" + return f"{prefix}{_SEPARATOR}v{FORMAT_VERSION}{_SEPARATOR}{name.hex()}" diff --git a/tensorrt_llm/_torch/disaggregation/blob/store.py b/tensorrt_llm/_torch/disaggregation/blob/store.py index eac8c10a9264..d2233e1aed75 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/store.py +++ b/tensorrt_llm/_torch/disaggregation/blob/store.py @@ -21,13 +21,22 @@ class StoreError(Exception): """A store-side failure the driver could not express as a row result.""" +class DriverRejected(StoreError): + """Raised by a driver on purpose, before it touched the store or the caller's memory. + + The only exception that certifies nothing escaped: the caller may forget the call. Any other + raise from ``register``, ``put`` or ``get`` may have escaped, and the caller keeps the memory + protected until the driver proves otherwise, at the latest through a ``drain`` returning True. + """ + + class RegistrationError(StoreError): - """Registering a region failed. + """Registering a region failed after the store was asked. Args: reason: What went wrong. handle: The handle to ``unregister`` later when the store may have registered the region - anyway, or ``None`` when nothing escaped. + anyway, or ``None`` when the store certainly did not. """ def __init__(self, reason: str, handle: RegionHandle | None = None) -> None: @@ -36,23 +45,37 @@ def __init__(self, reason: str, handle: RegionHandle | None = None) -> None: class RowStatus(enum.Enum): - """The logical state of one row within an operation.""" + """The logical state of one row within an operation. + + READY and FAULT are the only two that carry a verdict; MISS, NOT_READY and DECLINED say that + nothing moved and nothing went wrong, so the backend reports the row as not delivered without + a fault and they never count toward isolation. + """ PENDING = "pending" """Not concluded yet.""" READY = "ready" - """For a get: the bytes landed whole and the driver will not write the row again. For a put: - the object is readable whole under its key.""" + """For a get: the payload landed whole and intact, and the driver will not write the row again. + For a put: a complete object is readable under the key by any reader now, whoever wrote it. A + store's "already exists" counts only once that object is committed.""" MISS = "miss" - """For a get: the store holds no object under the key. Not a failure.""" + """For a get: the store holds no object under the key.""" + + 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.""" + + DECLINED = "declined" + """For a put: the store would not take the row, for space or policy. Nothing was published; + the other rows of the operation go on.""" CANCELLED = "cancelled" """Never started, because of a cancel or a drain; the row's memory was not touched.""" FAULT = "fault" - """Anything else: timeout, size or format mismatch, transport or store error.""" + """Anything else: timeout, size or integrity mismatch, transport or store error.""" @dataclass(frozen=True) @@ -74,6 +97,8 @@ def concluded(self) -> bool: 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) @@ -133,23 +158,21 @@ def nbytes(self) -> int: @dataclass(frozen=True) class RowIO: - """One row to move: its store key, where its payload sits and where its trailer sits. + """One row to move: its store key and where its payload sits. - The payload is the KV bytes of one staging slot; the trailer is the format's fixed tail, in a - buffer of the backend's own. The two spans never overlap. ``position`` is ``(group, row)`` - within the submission, kept so the backend can map results back; drivers treat it as opaque. + The payload is the KV bytes of one staging slot in a registered region; the object a driver + stores is that payload, in whatever format and with whatever integrity data the driver keeps. + ``position`` is ``(group, row)`` within the submission, kept so the backend can map results + back; drivers treat it as opaque. """ key: str payload: Span - trailer: Span position: tuple[int, int] | None = None def __post_init__(self) -> None: if not self.key: raise ValueError("a row needs a store key") - if self.payload.overlaps(self.trailer): - raise ValueError(f"payload and trailer overlap for key {self.key}") class RegionHandle(Protocol): @@ -165,63 +188,92 @@ class ProbeStatus(enum.Enum): HIT = "hit" MISS = "miss" + NOT_READY = "not_ready" UNKNOWN = "unknown" @dataclass(frozen=True) class DriverCapabilities: - """What a driver implements beyond the required methods. + """What a driver implements. Fixed for the driver's life; every field is declared. Args: + get: ``get`` is implemented. + put: ``put`` is implemented. cancel: ``cancel`` does something; without it the backend only suppresses unstarted work. probe: ``probe`` answers; without it the backend fetches blindly. max_rows_per_op: Rows one ``put``/``get`` accepts, or ``None`` for no limit. """ - cancel: bool = False - probe: bool = False - max_rows_per_op: int | None = None + get: bool + put: bool + cancel: bool + probe: bool + max_rows_per_op: int | None + + def __post_init__(self) -> None: + if not (self.get or self.put): + raise ValueError("a driver implements at least one of get and put") + if self.max_rows_per_op is not None and self.max_rows_per_op <= 0: + raise ValueError("max_rows_per_op is positive or None") @runtime_checkable class StoreDriver(Protocol): """What the backend asks of a store. - Threads: ``register``, ``unregister``, ``put``, ``get``, ``cancel`` and ``drain`` are called one - at a time from the backend's own threads; ``poll`` and ``access_ended`` may be called from any - thread, concurrently with the rest. A driver over a blocking SDK blocks inside ``put``/``get``; - the backend keeps those calls off the submission thread. - - Side effects: ``register``, ``put`` and ``get`` raise ``ValueError`` only before anything - escaped, so the caller may forget the call. Once ``put``/``get`` has returned, every later - problem is a row result, and the operation holds the rows' memory until ``access_ended`` is - true. A failure return is not access-end evidence. + Threads: ``register``, ``unregister``, ``drain`` and ``close`` are called one at a time. + ``put``, ``get``, ``cancel`` and ``probe`` may run concurrently with each other and with those, + from the backend's own threads; a blocked ``put``/``get`` must not delay another. ``poll``, + ``access_ended`` and ``access_unprovable`` may be called from any thread. Locks inside a driver + guard its state only and are never held while waiting on the store. + + Side effects: :class:`DriverRejected` is the only raise that certifies nothing escaped, and + ``put``/``get`` raise it only before recording the operation. Once recorded, the call returns + the handle and every later failure, of any kind, is a row result; the operation holds the + rows' memory until ``access_ended`` is true. Neither a failure result nor a successful + ``unregister`` is access-end evidence. + + Integrity: a READY get row is the whole object, intact, under the right key; a READY put row + is a complete object any reader can get. How a driver assures this is its own: a store's + native size checks, hidden unfinished objects and checksums, or data the driver keeps beside + the object. A failed, declined or cancelled put never damages an existing complete object. + + Queries (``poll``, ``access_ended``, ``access_unprovable``, ``cancel``) read recorded state: + they never block, never ask the store and never raise for a handle this driver issued. A handle + from elsewhere, or any call after ``close``, is a programming error and raises. """ @property - def capabilities(self) -> DriverCapabilities: ... + def capabilities(self) -> DriverCapabilities: + """The same value on every call.""" + ... def register(self, address: int, nbytes: int, layout: RegionLayout) -> RegionHandle: """Register one host region the driver may read and write. Raises: - ValueError: Bad geometry or overlap with a live registration; nothing escaped. + DriverRejected: Bad geometry or overlap with a live registration; nothing escaped. RegistrationError: The store refused; ``handle`` is set when the store may have registered the region anyway. """ ... def unregister(self, handle: RegionHandle) -> None: - """Revoke a registration. Idempotent once it succeeded; a raise leaves the registration - live and the call retryable.""" + """Revoke a registration. + + Called only once every operation touching the region has ``access_ended``; otherwise + raises :class:`DriverRejected` and the registration stays live. Idempotent once it + succeeded; any other raise leaves the registration live and the call retryable. + """ ... def put(self, rows: Sequence[RowIO]) -> OpHandle: """Publish ``rows``; may return before any byte moved. Raises: - ValueError: A span lies outside a live registration, or ``rows`` exceed the driver's - capabilities; nothing escaped. + DriverRejected: ``put`` is not declared, the driver is draining or closed, a key + appears twice in ``rows``, ``rows`` exceed ``max_rows_per_op``, or a payload lies + outside a live registration; nothing escaped. """ ... @@ -230,12 +282,26 @@ def get(self, rows: Sequence[RowIO]) -> OpHandle: ... def poll(self, op: OpHandle) -> tuple[RowResult, ...]: - """One result per row, in submission order; non-blocking. A concluded row never changes.""" + """One result per row, in submission order. A concluded row never changes.""" ... def access_ended(self, op: OpHandle) -> bool: """True once the operation can no longer touch any of its rows' memory, with evidence. - Never turns false again. Independent of ``poll``.""" + + Never turns false again. Independent of ``poll``. Without evidence the answer is False: + a timeout, an error or a guess never turns it true. + """ + ... + + def access_unprovable(self, op: OpHandle) -> bool: + """True once the driver's own mechanisms can never prove this operation's access ended, + for example when the store failed before issuing anything the driver could wait on. + + Never turns false again. Not permission to reuse memory: the rows stay protected until + evidence arrives from elsewhere, for example the store's own teardown; the backend counts + such operations instead of waiting for them. A driver that cannot tell always answers + False. + """ ... def cancel(self, op: OpHandle) -> None: @@ -244,13 +310,30 @@ def cancel(self, op: OpHandle) -> None: ... def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...]: - """Advisory presence of ``keys``, within ``budget_s`` seconds; UNKNOWN when the budget runs - out. Called only when ``capabilities.probe`` is set. Raises on a store error rather than - answering MISS.""" + """Advisory presence of ``keys``, answered within ``budget_s`` seconds. + + 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. Raises on a store error rather than answering MISS. + """ ... def drain(self, budget_s: float) -> bool: """Stop starting work and wait up to ``budget_s`` seconds for every operation to end its - access. True when none can still touch memory; False means keep everything registered and - try again. Never forces an end.""" + access and for the driver's own background work, such as outstanding probes, to finish. + + One-way: after the first call the driver rejects new ``put``/``get`` for good, whatever + the result. Returns False at once when an operation that cannot prove its end remains. + True when nothing can still touch memory and nothing is left running; False means keep + everything registered and try again. Never forces an end. + """ + ... + + def close(self) -> None: + """Release the store client and every resource the driver holds. + + Allowed only after ``drain`` returned True and every registration is unregistered; + otherwise raises :class:`DriverRejected` and changes nothing. Idempotent. + """ ... diff --git a/tests/unittest/_torch/disaggregation/blob/driver_conformance.py b/tests/unittest/_torch/disaggregation/blob/driver_conformance.py new file mode 100644 index 000000000000..75275264886a --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/driver_conformance.py @@ -0,0 +1,503 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""What any ``StoreDriver`` owes the backend, checked through the protocol alone. + +A driver's test module subclasses :class:`DriverConformance` as ``Test...`` and provides +``make_driver``. Every check observes consequences the backend could observe: bytes in registered +memory, row results, access end, rejections without effect. Nothing here uses a driver's own +controls, so the same checks run against the in-memory driver, a real store's driver, and the +deliberately broken drivers that prove each check bites. + +Each check runs inside a :class:`Harness` that keeps every endpoint and operation it created and, +when the check is over, waits for access to end, unregisters the regions, drains and closes the +driver. Memory whose access never provably ended is kept alive instead of being freed. +""" + +from __future__ import annotations + +import contextlib +import hashlib +import threading +import time +from collections.abc import Callable, Generator, Sequence + +import pytest +from host_memory import HostBuffer + +from tensorrt_llm._torch.disaggregation.blob.keys import encode_key +from tensorrt_llm._torch.disaggregation.blob.store import ( + DriverRejected, + OpHandle, + ProbeStatus, + RegionHandle, + RegionLayout, + RegistrationError, + RowIO, + RowResult, + RowStatus, + Span, + StoreDriver, + concluded, +) + +SLOT = 8 +PREFIX = "t" + +# Lender names are 54 bytes and start with a 16-byte scope digest shared by every name of a +# deployment. A key built from a truncated name collides across these. +SCOPE = b"scope" + bytes(11) +NAME_BYTES = 54 + + +def name_of(label: str) -> bytes: + return SCOPE + hashlib.blake2b(label.encode(), digest_size=NAME_BYTES - len(SCOPE)).digest() + + +NAMES = {label: name_of(label) for label in "ABC"} + + +def key_of(label: str) -> str: + return encode_key(PREFIX, NAMES[label]) + + +def wait_until(predicate: Callable[[], bool], what: str, timeout: float) -> None: + """Polls ``predicate`` until true; fails after ``timeout`` seconds.""" + deadline = time.monotonic() + timeout + while not predicate(): + assert time.monotonic() < deadline, f"timed out waiting for {what}" + time.sleep(0.001) + + +def statuses(driver: StoreDriver, op: OpHandle) -> list[RowStatus]: + return [result.status for result in driver.poll(op)] + + +class Endpoint: + """One side of a transfer: one registered region with a slot per label.""" + + def __init__( + self, + driver: StoreDriver, + labels: Sequence[str], + *, + fill: int = 0, + slot: int = SLOT, + register: bool = True, + ) -> None: + self.driver = driver + self.labels = list(labels) + self.slot = slot + self.payload = HostBuffer(slot * len(labels), fill) + self.handle: RegionHandle | None = None + if register: + self.register() + + def register(self) -> None: + """Registers the region. When the store may have kept it despite an error, the handle + is recorded all the same, so the region can still be unregistered later.""" + layout = RegionLayout(self.slot, len(self.labels), "kv") + try: + self.handle = self.driver.register(self.payload.address, self.payload.nbytes, layout) + except RegistrationError as error: + self.handle = error.handle + raise + + def rows(self) -> list[RowIO]: + return [ + RowIO(key_of(label), self.span(i), position=(0, i)) + for i, label in enumerate(self.labels) + ] + + def span(self, index: int) -> Span: + return self.payload.span(index * self.slot, self.slot) + + def fill_row(self, index: int, data: bytes) -> None: + self.payload.write(data, index * self.slot) + + def read_row(self, index: int) -> bytes: + return self.payload.read(index * self.slot, self.slot) + + def intact(self) -> bool: + return self.payload.canaries_intact() + + def unregister(self) -> None: + if self.handle is not None: + self.driver.unregister(self.handle) + + +# Endpoints whose memory a driver may still touch or still holds registered when a check ends. +# Never freed. +UNRELEASED_ENDPOINTS: list[Endpoint] = [] + + +class Harness: + """One driver under test plus everything a check created with it.""" + + def __init__(self, driver: StoreDriver, timeout: float) -> None: + self.driver = driver + self.timeout = timeout + self.endpoints: list[Endpoint] = [] + self.ops: list[OpHandle] = [] + self.done = False + + # Creating things the harness will clean up. + + def endpoint(self, labels: Sequence[str], *, fill: int = 0, slot: int = SLOT) -> Endpoint: + endpoint = Endpoint(self.driver, labels, fill=fill, slot=slot, register=False) + self.endpoints.append(endpoint) # kept before the store is asked, whatever register does + endpoint.register() + return endpoint + + def put(self, rows: Sequence[RowIO]) -> OpHandle: + op = self.driver.put(rows) + self.ops.append(op) + return op + + def get(self, rows: Sequence[RowIO]) -> OpHandle: + op = self.driver.get(rows) + self.ops.append(op) + return op + + def publish(self, rows: dict[str, bytes], slot: int = SLOT) -> Endpoint: + """Puts ``rows`` (label to payload) and returns the writer endpoint.""" + writer = self.endpoint(list(rows), slot=slot) + for i, data in enumerate(rows.values()): + writer.fill_row(i, data) + op = self.put(writer.rows()) + assert [r.status for r in self.settle(op)] == [RowStatus.READY] * len(rows) + self.settle_access(op) + return writer + + # Waiting. + + def settle(self, op: OpHandle) -> tuple[RowResult, ...]: + """Waits for every row's result, not for access end.""" + wait_until( + lambda: concluded(self.driver.poll(op)), "the operation to conclude", self.timeout + ) + return self.driver.poll(op) + + def settle_access(self, op: OpHandle) -> None: + wait_until(lambda: self.driver.access_ended(op), "access to end", self.timeout) + + def needs(self, *, put: bool = False, get: bool = False, rows: int = 1) -> None: + """Skips the check when the driver does not declare what it needs.""" + caps = self.driver.capabilities + if put and not caps.put: + pytest.skip("this driver declares no put") + if get and not caps.get: + pytest.skip("this driver declares no get") + if caps.max_rows_per_op is not None and rows > caps.max_rows_per_op: + pytest.skip(f"this driver takes at most {caps.max_rows_per_op} rows per operation") + + # Cleanup. + + def cleanup(self) -> list[str]: + """Ends everything the check created, in order; returns what did not go to plan. + + Memory a driver may still touch, or still holds registered, is retained for good. A + cleanup that did not finish may be called again. + """ + if self.done: + return [] + problems: list[str] = [] + for op in self.ops: + try: + self.settle_access(op) + except Exception as error: # noqa: BLE001 - every cleanup step is reported + problems.append(f"access end: {error!r}") + if problems: + UNRELEASED_ENDPOINTS.extend(self.endpoints) + return problems + for endpoint in self.endpoints: + try: + endpoint.unregister() + except Exception as error: # noqa: BLE001 - every cleanup step is reported + problems.append(f"unregister: {error!r}") + UNRELEASED_ENDPOINTS.append(endpoint) + if problems: + return problems + try: + if self.driver.drain(self.timeout): + self.driver.close() + else: + problems.append("drain did not return True") + except Exception as error: # noqa: BLE001 - every cleanup step is reported + problems.append(f"drain/close: {error!r}") + self.done = not problems + return problems + + +class DriverConformance: + """Contract checks for one driver. Subclass as ``Test...`` and implement ``make_driver``.""" + + timeout = 5.0 + + def make_driver(self) -> StoreDriver: + raise NotImplementedError + + def make_blocked_probe(self) -> tuple[StoreDriver, str, Callable[[], None]] | None: + """A driver whose store stalls every query of the returned key until the returned + callable runs; ``None`` when this driver cannot stage that, which skips the check.""" + return None + + @contextlib.contextmanager + def harness(self, driver: StoreDriver | None = None) -> Generator[Harness, None, None]: + """A fresh driver (or ``driver``) with cleanup afterwards. A failing check keeps its own + error; cleanup trouble fails only a check that passed.""" + h = Harness(driver if driver is not None else self.make_driver(), self.timeout) + try: + yield h + except BaseException: + h.cleanup() + raise + problems = h.cleanup() + assert not problems, "cleanup after the check: " + "; ".join(problems) + + # Bytes and results. + + def test_put_then_get_moves_whole_rows_and_never_aliases_the_source(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True, rows=2) + writer = h.publish({"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + writer.fill_row(0, b"xxxxxxxx") # the source changes after the put concluded + + reader = h.endpoint(["A", "B"], fill=0x2D) + results = h.settle(h.get(reader.rows())) + assert [r.status for r in results] == [RowStatus.READY, RowStatus.READY] + assert reader.read_row(0) == b"aaaaaaaa" + assert reader.read_row(1) == b"bbbbbbbb" + assert reader.intact() and writer.intact() + + def test_get_reports_hit_hit_miss_and_leaves_the_miss_untouched(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True, rows=3) + h.publish({"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + reader = h.endpoint(["A", "B", "C"], fill=0x2D) + results = h.settle(h.get(reader.rows())) + assert [r.status for r in results] == [ + RowStatus.READY, + RowStatus.READY, + RowStatus.MISS, + ] + assert reader.read_row(2) == b"-" * SLOT + assert reader.intact() + + def test_an_object_of_the_wrong_size_is_a_fault_not_a_miss(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True) + h.publish({"A": b"aaaa"}, slot=4) + reader = h.endpoint(["A"], fill=0x2D) + (result,) = h.settle(h.get(reader.rows())) + assert result.status is RowStatus.FAULT + assert reader.read_row(0) == b"-" * SLOT + assert reader.intact() + + def test_a_duplicate_put_of_a_committed_object_is_ready(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True) + h.publish({"A": b"aaaaaaaa"}) + h.publish({"A": b"aaaaaaaa"}) # a second writer, same content + reader = h.endpoint(["A"]) + (result,) = h.settle(h.get(reader.rows())) + assert result.status is RowStatus.READY + assert reader.read_row(0) == b"aaaaaaaa" + + def test_a_put_under_an_existing_key_leaves_the_object_whole(self) -> None: + """Whatever a driver answers a writer bringing other bytes, the object readers get is + the one that was committed.""" + with self.harness() as h: + h.needs(put=True, get=True) + h.publish({"A": b"aaaaaaaa"}) + other = h.endpoint(["A"]) + other.fill_row(0, b"zzzzzzzz") + (result,) = h.settle(h.put(other.rows())) + assert result.status in (RowStatus.READY, RowStatus.FAULT) + reader = h.endpoint(["A"], fill=0x2D) + (got,) = h.settle(h.get(reader.rows())) + assert got.status is RowStatus.READY + assert reader.read_row(0) == b"aaaaaaaa" + + # Rejections leave no trace. + + def test_rows_outside_a_live_registration_are_rejected_before_any_effect(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True) + inside = h.endpoint(["A"]) + stray = HostBuffer(SLOT, fill=0x2D) + h.publish({"A": b"aaaaaaaa"}) + + with pytest.raises(DriverRejected): + h.get([RowIO(key_of("A"), stray.span())]) + assert stray.read() == b"-" * SLOT and stray.canaries_intact() + + inside.fill_row(0, b"bbbbbbbb") + inside.unregister() + with pytest.raises(DriverRejected): + h.put([RowIO(key_of("B"), inside.span(0))]) + reader = h.endpoint(["B"]) + assert [r.status for r in h.settle(h.get(reader.rows()))] == [RowStatus.MISS] + + def test_a_key_appears_at_most_once_per_operation(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True, rows=2) + writer = h.endpoint(["A", "A"]) + writer.fill_row(0, b"aaaaaaaa") + writer.fill_row(1, b"bbbbbbbb") + with pytest.raises(DriverRejected): + h.put(writer.rows()) + reader = h.endpoint(["A", "A"], fill=0x2D) + with pytest.raises(DriverRejected): + h.get(reader.rows()) + single = h.endpoint(["A"]) + assert [r.status for r in h.settle(h.get(single.rows()))] == [RowStatus.MISS] + assert reader.read_row(0) == reader.read_row(1) == b"-" * SLOT + + def test_rows_beyond_the_declared_limit_are_rejected(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True) + limit = h.driver.capabilities.max_rows_per_op + if limit is None or limit >= 3: + pytest.skip("this driver declares no row limit this check can exceed") + writer = h.endpoint(["A", "B", "C"][: limit + 1]) + with pytest.raises(DriverRejected): + h.put(writer.rows()) + reader = h.endpoint(["A"]) + assert [r.status for r in h.settle(h.get(reader.rows()))] == [RowStatus.MISS] + + def test_an_undeclared_direction_is_rejected(self) -> None: + with self.harness() as h: + caps = h.driver.capabilities + if caps.get and caps.put: + pytest.skip("this driver declares both directions") + endpoint = h.endpoint(["A"]) + endpoint.fill_row(0, b"aaaaaaaa") + with pytest.raises(DriverRejected): + if caps.put: + h.get(endpoint.rows()) + else: + h.put(endpoint.rows()) + assert endpoint.read_row(0) == b"aaaaaaaa" + + # Results, access and capabilities over time. + + def test_results_are_stable_and_access_ends_with_evidence(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True, rows=2) + h.publish({"A": b"aaaaaaaa"}) + reader = h.endpoint(["A", "C"]) + op = h.get(reader.rows()) + results = h.settle(op) + assert h.driver.poll(op) == results + h.settle_access(op) + assert h.driver.access_ended(op) + assert not h.driver.access_unprovable(op) + assert h.driver.poll(op) == results + + def test_cancel_after_conclusion_changes_nothing(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True) + if not h.driver.capabilities.cancel: + pytest.skip("this driver declares no cancel") + h.publish({"A": b"aaaaaaaa"}) + reader = h.endpoint(["A"], fill=0x2D) + op = h.get(reader.rows()) + results = h.settle(op) + h.driver.cancel(op) + assert h.driver.poll(op) == results + assert reader.read_row(0) == b"aaaaaaaa" + + def test_capabilities_do_not_change(self) -> None: + with self.harness() as h: + h.needs(put=True) + before = h.driver.capabilities + h.publish({"A": b"aaaaaaaa"}) + assert h.driver.capabilities == before + assert h.driver.capabilities == before + + def test_probe_reports_hit_and_miss(self) -> None: + with self.harness() as h: + h.needs(put=True) + if not h.driver.capabilities.probe: + pytest.skip("this driver declares no probe") + h.publish({"A": b"aaaaaaaa"}) + assert h.driver.probe([key_of("A"), key_of("C")], self.timeout) == ( + ProbeStatus.HIT, + ProbeStatus.MISS, + ) + + def test_probe_returns_within_its_budget(self) -> None: + with self.harness() as h: + if not h.driver.capabilities.probe: + pytest.skip("this driver declares no probe") + keys = [key_of(label) for label in "ABC"] + started = time.monotonic() + answer = h.driver.probe(keys, 0.05) + elapsed = time.monotonic() - started + assert elapsed < 0.05 + 0.5, f"probe took {elapsed:.2f}s on a 0.05s budget" + assert len(answer) == 3 and set(answer) <= set(ProbeStatus) + + def test_probe_answers_unknown_for_a_key_the_store_does_not_answer_in_time(self) -> None: + staged = self.make_blocked_probe() + if staged is None: + pytest.skip("this driver cannot stage a blocked query") + driver, blocked_key, unblock = staged + with self.harness(driver) as h: + if not driver.capabilities.probe: + pytest.skip("this driver declares no probe") + answer: list[tuple[ProbeStatus, ...]] = [] + # Off this thread, so a probe that waits for the store fails instead of hanging. + prober = threading.Thread( + target=lambda: answer.append(driver.probe([blocked_key], 0.05)), daemon=True + ) + prober.start() + prober.join(0.05 + 0.5) + try: + assert not prober.is_alive(), "probe waited for the store past its budget" + assert answer == [(ProbeStatus.UNKNOWN,)] + finally: + unblock() + prober.join(h.timeout) + + # Registration, drain and close. + + def test_unregister_is_idempotent_and_frees_the_range(self) -> None: + with self.harness() as h: + endpoint = h.endpoint(["A"]) + with pytest.raises(DriverRejected): + h.driver.register(endpoint.payload.address, SLOT, RegionLayout(SLOT, 1)) + endpoint.unregister() + endpoint.unregister() + try: + handle = h.driver.register(endpoint.payload.address, SLOT, RegionLayout(SLOT, 1)) + except DriverRejected as error: + pytest.fail(f"the range stayed registered: {error}") + h.driver.unregister(handle) + + def test_drain_is_one_way(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True) + writer = h.publish({"A": b"aaaaaaaa"}) + assert h.driver.drain(self.timeout) is True + with pytest.raises(DriverRejected): + h.put(writer.rows()) + with pytest.raises(DriverRejected): + h.get(writer.rows()) + assert h.driver.drain(self.timeout) is True + assert writer.read_row(0) == b"aaaaaaaa" + + def test_close_needs_a_true_drain_and_no_live_registrations(self) -> None: + with self.harness() as h: + h.needs(put=True) + writer = h.publish({"A": b"aaaaaaaa"}) + with pytest.raises(DriverRejected): + h.driver.close() + assert h.driver.drain(self.timeout) is True + with pytest.raises(DriverRejected): + h.driver.close() + writer.unregister() + h.driver.close() + h.driver.close() + h.done = True # closed by the check itself + with pytest.raises(DriverRejected): + h.put(writer.rows()) diff --git a/tests/unittest/_torch/disaggregation/blob/memory_store.py b/tests/unittest/_torch/disaggregation/blob/memory_store.py index eab2f9cf48c1..1fb0c94bb2c7 100644 --- a/tests/unittest/_torch/disaggregation/blob/memory_store.py +++ b/tests/unittest/_torch/disaggregation/blob/memory_store.py @@ -2,31 +2,41 @@ # SPDX-License-Identifier: Apache-2.0 """An in-process store and a driver over it, for tests. -``MemoryStore`` is a dictionary of complete objects that two endpoints in one process can share. -``MemoryDriver`` really copies bytes between registered host memory and that store. In automatic -mode an operation completes inside ``put``/``get``, like a blocking SDK. In controlled mode the test -releases three things separately per operation: the transfer, the logical results and the end of -access, so it can stage partial arrival, a result before or after access end, a late write after a -fault, and registration failures. No threads are started here. Transfers copy and call the store -outside the driver's lock, so queries never wait on the store; the lock is held around state -changes and around a late write's short copy, which keeps ``end_access`` out until it lands. +``MemoryStore`` is a dictionary of complete objects that endpoints in one process can share. It +can be given a capacity, and a test can mark a key as held by an unfinished write. ``MemoryDriver`` +really copies bytes between registered host memory and that store, keeping a trailer beside each +object that it verifies on every get, the way a real store checks sizes and checksums natively. + +In automatic mode an operation completes inside ``put``/``get``, like a blocking SDK. In controlled +mode the test releases three things separately per operation: the transfer, the logical results +and the end of access, so it can stage partial arrival, a result before or after access end, a late +write after a fault, an operation whose end can never be proven, and registration failures. + +Transfers copy and call the store outside the driver's lock, so queries never wait on the store +and one blocked operation never delays another. Each submission check is its own method, so a +test can subclass the driver and disable exactly one rule. """ from __future__ import annotations import ctypes +import enum import threading from collections.abc import Iterable, Sequence from dataclasses import dataclass from typing import Literal -from tensorrt_llm._torch.disaggregation.blob.format import TRAILER_BYTES +from object_trailer import Checksum, check_trailer, encode_trailer + from tensorrt_llm._torch.disaggregation.blob.store import ( CANCELLED, + DECLINED, MISS, + NOT_READY, PENDING, READY, DriverCapabilities, + DriverRejected, OpHandle, ProbeStatus, RegionHandle, @@ -40,6 +50,35 @@ faulted, ) +DEFAULT_CAPABILITIES = DriverCapabilities( + get=True, put=True, cancel=True, probe=True, max_rows_per_op=None +) +MAX_PROBES_IN_FLIGHT = 4 + +Kind = Literal["put", "get"] + + +class ObjectNotReady(StoreError): + """The key is held by a write that has not finished.""" + + +class Commit(enum.Enum): + """How the store answered a commit.""" + + STORED = "stored" + SAME = "same" # an identical object was already there + DIFFERENT = "different" # a different object holds the key; it was kept + DECLINED = "declined" # no room + NOT_READY = "not_ready" # another write holds the key + + +@dataclass(frozen=True) +class StoredObject: + """One committed object: its payload and the trailer the driver keeps beside it.""" + + payload: bytes + trailer: bytes + class MemoryStore: """Objects by key. Each object is a private copy, committed whole or not at all. @@ -48,43 +87,78 @@ class MemoryStore: raise ``StoreError``, which the driver turns into a row fault. """ - def __init__(self) -> None: + def __init__(self, capacity: int | None = None) -> None: self._lock = threading.Lock() - self._objects: dict[str, bytes] = {} + self._objects: dict[str, StoredObject] = {} + self._writing: set[str] = set() + self.capacity = capacity - def commit(self, key: str, data: bytes) -> bool: - """Store ``data`` under ``key``. False when a different object already holds the key; - the existing object is kept. Storing identical bytes again is a no-op success.""" - copy = bytes(data) + def commit(self, key: str, payload: bytes, trailer: bytes) -> Commit: + new = StoredObject(bytes(payload), bytes(trailer)) with self._lock: + if key in self._writing: + return Commit.NOT_READY existing = self._objects.get(key) - if existing is None: - self._objects[key] = copy - return True - return existing == copy - - def get(self, key: str) -> bytes | None: + if existing is not None: + return Commit.SAME if existing == new else Commit.DIFFERENT + if self.capacity is not None and len(self._objects) >= self.capacity: + return Commit.DECLINED + self._objects[key] = new + return Commit.STORED + + def get(self, key: str) -> StoredObject | None: with self._lock: + if key in self._writing: + raise ObjectNotReady(key) return self._objects.get(key) + def presence(self, key: str) -> ProbeStatus: + with self._lock: + if key in self._writing: + return ProbeStatus.NOT_READY + return ProbeStatus.HIT if key in self._objects else ProbeStatus.MISS + def contains(self, key: str) -> bool: - return self.get(key) is not None + return self.presence(key) is ProbeStatus.HIT - def committed_object(self, key: str) -> bytes: - """The whole object under ``key``; ``KeyError`` when nothing is committed there.""" - data = self.get(key) - if data is None: + def committed_object(self, key: str) -> StoredObject: + """The object under ``key``; ``KeyError`` when nothing is committed there.""" + stored = self.get(key) + if stored is None: raise KeyError(f"no committed object under {key!r}") - return data + return stored - def committed_payload(self, key: str, trailer_nbytes: int = TRAILER_BYTES) -> bytes: - """The payload of the object under ``key``, without its trailer.""" - return self.committed_object(key)[:-trailer_nbytes] + def committed_payload(self, key: str) -> bytes: + return self.committed_object(key).payload def __len__(self) -> int: with self._lock: return len(self._objects) + # Test controls. + + def begin_write(self, key: str) -> None: + """Hold ``key`` as another writer's unfinished object until ``end_write``.""" + with self._lock: + self._writing.add(key) + + def end_write(self, key: str) -> None: + with self._lock: + self._writing.discard(key) + + def corrupt(self, key: str, index: int) -> None: + """Flip one payload byte of the committed object, as a damaged store would.""" + with self._lock: + stored = self._objects[key] + payload = bytearray(stored.payload) + payload[index] ^= 0xFF + self._objects[key] = StoredObject(bytes(payload), stored.trailer) + + def overwrite(self, key: str, payload: bytes, trailer: bytes) -> None: + """Replace whatever holds ``key``. Only a deliberately broken driver calls this.""" + with self._lock: + self._objects[key] = StoredObject(bytes(payload), bytes(trailer)) + @dataclass(eq=False) class _Registration: @@ -106,9 +180,16 @@ class _Row: class MemoryOp: """One put or get. Tests hold it to drive the controlled mode.""" - kind: Literal["put", "get"] + kind: Kind rows: list[_Row] access_ended: bool = False + unprovable: bool = False + + +@dataclass(eq=False) +class _Probe: + worker: threading.Thread + finished: threading.Event class MemoryDriver: @@ -119,7 +200,7 @@ def __init__( store: MemoryStore, *, controlled: bool = False, - capabilities: DriverCapabilities = DriverCapabilities(cancel=True, probe=True), + capabilities: DriverCapabilities = DEFAULT_CAPABILITIES, ) -> None: self.store = store self.controlled = controlled @@ -127,7 +208,9 @@ def __init__( self._lock = threading.Condition(threading.RLock()) self._registrations: list[_Registration] = [] self._ops: list[MemoryOp] = [] + self._probes: list[_Probe] = [] self._draining = False + self._closed = False self._cancel_effective = True self._probe_unknown = False self._register_failure: tuple[str, bool] | None = None @@ -140,13 +223,17 @@ def capabilities(self) -> DriverCapabilities: return self._capabilities def register(self, address: int, nbytes: int, layout: RegionLayout) -> _Registration: - span = Span(address, nbytes) + try: + span = Span(address, nbytes) + except ValueError as error: + raise DriverRejected(str(error)) from None if layout.nbytes > nbytes: - raise ValueError(f"layout needs {layout.nbytes} bytes, region has {nbytes}") + raise DriverRejected(f"layout needs {layout.nbytes} bytes, region has {nbytes}") with self._lock: + self._require_accepting() for other in self._registrations: if other.live and other.span.overlaps(span): - raise ValueError(f"region {span} overlaps a live registration {other.span}") + raise DriverRejected(f"region {span} overlaps a live registration {other.span}") registration = _Registration(span, layout) if self._register_failure is not None: reason, escaped = self._register_failure @@ -162,6 +249,9 @@ def register(self, address: int, nbytes: int, layout: RegionLayout) -> _Registra def unregister(self, handle: RegionHandle) -> None: with self._lock: registration = self._registration(handle) + if not registration.live: + return + self._check_region_idle(registration) reason = self._unregister_failures.pop(id(registration), None) if reason is not None: raise StoreError(reason) @@ -181,33 +271,72 @@ def access_ended(self, op: OpHandle) -> bool: with self._lock: return self._op(op).access_ended + def access_unprovable(self, op: OpHandle) -> bool: + with self._lock: + return self._op(op).unprovable + def cancel(self, op: OpHandle) -> None: with self._lock: if self._cancel_effective: self._cancel_unstarted(self._op(op)) def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...]: - if self._probe_unknown: - return tuple(ProbeStatus.UNKNOWN for _ in keys) - return tuple( - ProbeStatus.HIT if self.store.contains(key) else ProbeStatus.MISS for key in keys - ) + answers: dict[str, ProbeStatus] = {} + errors: list[Exception] = [] + finished = threading.Event() + + def ask() -> None: + try: + for key in keys: + answers[key] = self.store.presence(key) + except Exception as error: # reported from the caller's thread below + errors.append(error) + finally: + finished.set() + with self._lock: + self._lock.notify_all() + + # The store may block; the budget bounds this call, not the store. Registration and + # start happen under the lock, so close and drain see every worker. + with self._lock: + self._require_not_closed() + self._probes = [probe for probe in self._probes if not probe.finished.is_set()] + if self._draining or self._probe_unknown or len(self._probes) >= MAX_PROBES_IN_FLIGHT: + return tuple(ProbeStatus.UNKNOWN for _ in keys) + worker = threading.Thread(target=ask, name="memory-driver-probe", daemon=True) + self._probes.append(_Probe(worker, finished)) + worker.start() + if finished.wait(budget_s): + worker.join() # it is exiting; leave no thread behind a call that got its answer + if errors: + raise errors[0] + return tuple(answers.get(key, ProbeStatus.UNKNOWN) for key in keys) def drain(self, budget_s: float) -> bool: with self._lock: - self._draining = True + self._require_not_closed() + self._start_draining() for op in self._ops: self._cancel_unstarted(op) - return self._lock.wait_for( - lambda: all(op.access_ended for op in self._ops), timeout=budget_s - ) + self._lock.wait_for(self._drain_can_return, timeout=budget_s) + return self._quiet() + + def close(self) -> None: + with self._lock: + if self._closed: + return + self._check_closable() + self._closed = True + probes, self._probes = self._probes, [] + for probe in probes: # all finished; join only lets the threads exit + probe.worker.join() # Controlled mode: the test releases each stage. def transfer(self, op: MemoryOp, rows: Iterable[int] | None = None) -> None: """Move the bytes of ``rows`` (all by default) and stage their results. - The copies and store calls run outside the driver's lock; a ``StoreError`` from the store + The copies and store calls run outside the driver's lock; any raise from the store becomes that row's fault instead of escaping. """ with self._lock: @@ -241,7 +370,8 @@ def complete(self, op: MemoryOp, rows: Iterable[int] | None = None) -> None: self.publish_results(op, indexes) def fail(self, op: MemoryOp, reason: str, rows: Iterable[int] | None = None) -> None: - """Conclude ``rows`` as FAULT without moving bytes; access stays open.""" + """Conclude ``rows`` as FAULT without moving bytes, as a store that reported an error + would; access stays open because the store may still be writing.""" with self._lock: self._require_open(op) for index in self._indexes(op, rows): @@ -252,8 +382,8 @@ def fail(self, op: MemoryOp, reason: str, rows: Iterable[int] | None = None) -> def write_late(self, op: MemoryOp, index: int, data: bytes) -> None: """Write into a faulted get row's payload after its result, as a transfer that kept going - would. Refused once access has ended, and for rows that are READY, MISS, CANCELLED or - still pending: those must stay as reported.""" + would. Refused once access has ended, and for rows that are not FAULT: those must stay + as reported.""" with self._lock: self._require_open(op) if op.kind != "get": @@ -276,6 +406,14 @@ def end_access(self, op: MemoryOp) -> None: op.access_ended = True self._lock.notify_all() + def mark_unprovable(self, op: MemoryOp) -> None: + """Declare that the driver will never learn when this operation's access ends, as after a + store failure that issued nothing to wait on. Access stays open.""" + with self._lock: + self._require_open(op) + op.unprovable = True + self._lock.notify_all() + def set_cancel_effective(self, effective: bool) -> None: self._cancel_effective = effective @@ -298,72 +436,193 @@ def open_ops(self) -> list[MemoryOp]: with self._lock: return [op for op in self._ops if not op.access_ended] - # Internals. + def probes_in_flight(self) -> int: + with self._lock: + return sum(1 for probe in self._probes if not probe.finished.is_set()) + + def join_probes(self, timeout: float) -> None: + """Wait for every probe worker to exit, for tests that opened a gate they had closed.""" + with self._lock: + probes = list(self._probes) + for probe in probes: + probe.worker.join(timeout) + + # Submission checks, one rule each. + + def _check_direction(self, kind: Kind) -> None: + if not getattr(self._capabilities, kind): + raise DriverRejected(f"{kind} is not declared by this driver") + + def _check_not_draining(self) -> None: + if self._draining: + raise DriverRejected("the driver is draining") - def _submit(self, kind: Literal["put", "get"], rows: Sequence[RowIO]) -> MemoryOp: + def _check_row_limit(self, rows: Sequence[RowIO]) -> None: if not rows: - raise ValueError("an operation needs at least one row") + raise DriverRejected("an operation needs at least one row") limit = self._capabilities.max_rows_per_op if limit is not None and len(rows) > limit: - raise ValueError(f"{len(rows)} rows exceed the limit of {limit} per operation") + raise DriverRejected(f"{len(rows)} rows exceed the limit of {limit} per operation") + + def _check_unique_keys(self, rows: Sequence[RowIO]) -> None: + keys = [io.key for io in rows] + if len(set(keys)) != len(keys): + raise DriverRejected("a key appears more than once in the operation") + + def _check_registered(self, rows: Sequence[RowIO]) -> None: + for io in rows: + if not self._registered(io.payload): + raise DriverRejected(f"{io.payload} is outside every live registration") + + def _check_region_idle(self, registration: _Registration) -> None: + busy = sum( + 1 + for op in self._ops + if not op.access_ended + and any(row.io.payload.within(registration.span) for row in op.rows) + ) + if busy: + raise DriverRejected(f"{busy} operation(s) may still access the region") + + def _check_closable(self) -> None: + if not self._draining or not self._quiet(): + 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") + + def _start_draining(self) -> None: + self._draining = True + + def _drain_can_return(self) -> bool: + return self._quiet() or any(op.unprovable and not op.access_ended for op in self._ops) + + def _quiet(self) -> bool: + """Nothing can still touch memory and no background work is running.""" + return all(op.access_ended for op in self._ops) and all( + probe.finished.is_set() for probe in self._probes + ) + + # Internals. + + def _submit(self, kind: Kind, rows: Sequence[RowIO]) -> MemoryOp: with self._lock: - if self._draining: - raise ValueError("the driver is draining") - for io in rows: - for span in (io.payload, io.trailer): - if not self._registered(span): - raise ValueError(f"{span} is outside every live registration") + self._require_accepting() + self._check_direction(kind) + self._check_not_draining() + self._check_row_limit(rows) + self._check_unique_keys(rows) + self._check_registered(rows) op = MemoryOp(kind, [_Row(io) for io in rows]) self._ops.append(op) if not self.controlled: - # Like a blocking SDK: the call returns once every row is decided and access is over. - self.transfer(op) - self.publish_results(op) - self.end_access(op) + self._before_execute(op) + self._execute(op) return op - def _move(self, kind: str, io: RowIO) -> RowResult: + def _before_execute(self, op: MemoryOp) -> None: + """Between recording an operation and executing it; a test races a drain here.""" + + def _execute(self, op: MemoryOp) -> None: + """Automatic mode: like a blocking SDK, the call returns once every row is decided and + access is over. Rows a concurrent drain cancelled in between stay as they are.""" + with self._lock: + taken = [row for row in op.rows if not row.started and not row.published.concluded] + for row in taken: + row.started = True + for row in taken: + result = self._move(op.kind, row.io) + with self._lock: + row.published = result + row.transferred = True + with self._lock: + if not op.access_ended: + op.access_ended = True + self._lock.notify_all() + + def _move(self, kind: Kind, io: RowIO) -> RowResult: + # The operation is recorded by now: nothing may escape as an exception any more. try: return self._put_row(io) if kind == "put" else self._get_row(io) except StoreError as error: return faulted(str(error)) + except Exception as error: + return faulted(f"{type(error).__name__}: {error}") def _put_row(self, io: RowIO) -> RowResult: - data = ctypes.string_at(io.payload.address, io.payload.nbytes) - data += ctypes.string_at(io.trailer.address, io.trailer.nbytes) - if self.store.commit(io.key, data): + payload = ctypes.string_at(io.payload.address, io.payload.nbytes) + trailer = self._trailer_for(io, payload) + outcome = self.store.commit(io.key, payload, trailer) + if outcome in (Commit.STORED, Commit.SAME): return READY + if outcome is Commit.DECLINED: + return DECLINED + if outcome is Commit.NOT_READY: + return NOT_READY + # A diagnostic this store can afford; the contract asks no driver to compare bytes. return faulted("key exists with different content") def _get_row(self, io: RowIO) -> RowResult: - data = self.store.get(io.key) - if data is None: + try: + stored = self.store.get(io.key) + except ObjectNotReady: + return NOT_READY + if stored is None: return MISS - expected = io.payload.nbytes + io.trailer.nbytes - if len(data) != expected: - return faulted(f"object is {len(data)} bytes, expected {expected}") - ctypes.memmove(io.payload.address, data, io.payload.nbytes) - ctypes.memmove(io.trailer.address, data[io.payload.nbytes :], io.trailer.nbytes) + if len(stored.payload) != io.payload.nbytes: + return faulted(f"object is {len(stored.payload)} bytes, expected {io.payload.nbytes}") + problem = self._verify(io, stored) + if problem is not None: + return faulted(problem) + ctypes.memmove(io.payload.address, stored.payload, len(stored.payload)) return READY + def _trailer_for(self, io: RowIO, payload: bytes) -> bytes: + return encode_trailer( + io.key.encode(), io.payload.nbytes, checksum=Checksum.CRC32, payload=payload + ) + + def _verify(self, io: RowIO, stored: StoredObject) -> str | None: + """Why the stored object may not be delivered as ``io``, or ``None``.""" + return check_trailer(stored.trailer, io.key.encode(), io.payload.nbytes, stored.payload) + def _cancel_unstarted(self, op: MemoryOp) -> None: for row in op.rows: if not row.started and not row.published.concluded: row.published = CANCELLED + self._settle_untouched(op) + + def _settle_untouched(self, op: MemoryOp) -> None: + """An operation whose rows were all cancelled before any started never touched memory, + which is access-end evidence in itself.""" + if not op.access_ended and all( + row.published.status is RowStatus.CANCELLED and not row.started for row in op.rows + ): + op.access_ended = True + self._lock.notify_all() def _registered(self, span: Span) -> bool: return any(r.live and span.within(r.span) for r in self._registrations) def _registration(self, handle: RegionHandle) -> _Registration: + self._require_not_closed() if not isinstance(handle, _Registration) or handle not in self._registrations: raise ValueError("handle does not belong to this driver") return handle def _op(self, op: OpHandle) -> MemoryOp: + self._require_not_closed() if not isinstance(op, MemoryOp) or op not in self._ops: raise ValueError("operation does not belong to this driver") return op + def _require_accepting(self) -> None: + if self._closed: + raise DriverRejected("the driver is closed") + + def _require_not_closed(self) -> None: + if self._closed: + raise RuntimeError("the driver is closed") + def _require_open(self, op: MemoryOp) -> None: if self._op(op).access_ended: raise RuntimeError("the operation already ended its access") diff --git a/tests/unittest/_torch/disaggregation/blob/object_trailer.py b/tests/unittest/_torch/disaggregation/blob/object_trailer.py new file mode 100644 index 000000000000..2da2d42e679f --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/object_trailer.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A fixed-size trailer the in-memory driver keeps beside each object. + +The driver boundary leaves object format and integrity to each driver. A real store checks sizes, +hides unfinished objects and can checksum natively; the in-memory driver stands in for that with +this trailer, which it verifies on every get before a row counts as READY: version, key digest, +payload length and, when an algorithm is set, a payload checksum. +""" + +from __future__ import annotations + +import enum +import hashlib +import struct +import zlib + +TRAILER_VERSION = 1 +MAGIC = b"TLBK" + + +class Checksum(enum.IntEnum): + """Payload checksum algorithms a trailer can name.""" + + NONE = 0 + CRC32 = 1 + + +# magic, trailer version, checksum algorithm, name digest, payload length, payload checksum. +_TRAILER = struct.Struct("<4sHH16sQQ") +TRAILER_BYTES = _TRAILER.size +_DIGEST_BYTES = 16 + + +def name_digest(name: bytes) -> bytes: + """The 16-byte digest of an object name that the trailer carries.""" + return hashlib.blake2b(name, digest_size=_DIGEST_BYTES).digest() + + +def payload_checksum(algorithm: Checksum, payload: bytes | memoryview) -> int: + """The checksum of ``payload`` under ``algorithm``; 0 for ``Checksum.NONE``.""" + if algorithm is Checksum.NONE: + return 0 + return zlib.crc32(payload) & 0xFFFFFFFF + + +def encode_trailer( + name: bytes, + payload_nbytes: int, + *, + checksum: Checksum = Checksum.NONE, + payload: bytes | memoryview | None = None, +) -> bytes: + """The trailer for an object of ``payload_nbytes`` bytes stored under ``name``. + + Args: + name: The object's name; the in-memory driver uses the store key. + payload_nbytes: The payload length. + checksum: The algorithm to checksum the payload with. + payload: The payload bytes; required when ``checksum`` is not ``NONE``. + """ + if payload_nbytes <= 0: + raise ValueError(f"payload length must be positive, got {payload_nbytes}") + value = 0 + if checksum is not Checksum.NONE: + if payload is None or len(payload) != payload_nbytes: + raise ValueError("a checksum needs the whole payload") + value = payload_checksum(checksum, payload) + return _TRAILER.pack(MAGIC, TRAILER_VERSION, checksum, name_digest(name), payload_nbytes, value) + + +def check_trailer( + trailer: bytes, name: bytes, payload_nbytes: int, payload: bytes | memoryview | None = None +) -> str | None: + """``None`` when ``trailer`` matches ``name``, ``payload_nbytes`` and, if it names an algorithm, + the checksum of ``payload``; otherwise why it does not. Every trailer byte is checked.""" + if len(trailer) != TRAILER_BYTES: + return f"trailer is {len(trailer)} bytes, expected {TRAILER_BYTES}" + magic, version, algorithm, digest, length, value = _TRAILER.unpack(trailer) + if magic != MAGIC: + return "bad magic" + if version != TRAILER_VERSION: + return f"trailer version {version}, expected {TRAILER_VERSION}" + try: + checksum = Checksum(algorithm) + except ValueError: + return f"unknown checksum algorithm {algorithm}" + if digest != name_digest(name): + return "object name mismatch" + if length != payload_nbytes: + return f"payload length {length}, expected {payload_nbytes}" + if checksum is Checksum.NONE: + return None if value == 0 else "checksum set without an algorithm" + if payload is None: + return "checksum needs the payload" + if value != payload_checksum(checksum, payload): + return "payload checksum mismatch" + return None diff --git a/tests/unittest/_torch/disaggregation/blob/test_blob_format.py b/tests/unittest/_torch/disaggregation/blob/test_blob_format.py new file mode 100644 index 000000000000..deb724150bf3 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_blob_format.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The key encoding, the driver boundary's value types, and the in-memory driver's trailer.""" + +from __future__ import annotations + +import pytest +from driver_conformance import NAMES, SCOPE +from object_trailer import TRAILER_BYTES, TRAILER_VERSION, Checksum, check_trailer, encode_trailer + +from tensorrt_llm._torch.disaggregation.blob.keys import FORMAT_VERSION, encode_key +from tensorrt_llm._torch.disaggregation.blob.store import ( + DriverCapabilities, + RowIO, + RowResult, + RowStatus, + Span, + faulted, +) + +pytestmark = pytest.mark.cpu_only + + +# Keys. + + +def test_keys_use_the_whole_name_and_round_trip_through_hex() -> None: + keys = {label: encode_key("prod", name) for label, name in NAMES.items()} + assert len(set(keys.values())) == len(NAMES) + for label, key in keys.items(): + prefix, version, hexed = key.split("/") + assert (prefix, version) == ("prod", f"v{FORMAT_VERSION}") + assert bytes.fromhex(hexed) == NAMES[label] + assert NAMES[label].startswith(SCOPE) and len(NAMES[label]) == 54 + + +def test_key_prefix_is_a_store_partition_with_a_narrow_charset() -> None: + for prefix in ("prod", "a_b-1", "A" * 64): + assert encode_key(prefix, b"\x01").startswith(prefix + "/") + for prefix in ("", "a/b", "a b", "a.b", "A" * 65, "中"): + with pytest.raises(ValueError): + encode_key(prefix, b"\x01") + with pytest.raises(ValueError): + encode_key("prod", b"") + + +# Value types. + + +def test_row_result_pairs_fault_text_with_fault() -> None: + assert faulted("timeout").fault == "timeout" + assert RowResult(RowStatus.READY).concluded and not RowResult(RowStatus.PENDING).concluded + with pytest.raises(ValueError): + RowResult(RowStatus.FAULT) + with pytest.raises(ValueError): + RowResult(RowStatus.READY, "text") + + +def test_span_layout_and_row_validate_their_numbers() -> None: + with pytest.raises(ValueError): + Span(-1, 8) + with pytest.raises(ValueError): + Span(0, 0) + assert Span(0x1000, 64).overlaps(Span(0x1020, 64)) + assert Span(0x1010, 16).within(Span(0x1000, 64)) + with pytest.raises(ValueError): + RowIO("", Span(0x1000, 8)) + + +def test_capabilities_are_declared_in_full() -> None: + with pytest.raises(TypeError): + DriverCapabilities(get=True, put=True, cancel=False, probe=False) # type: ignore[call-arg] + with pytest.raises(ValueError): + DriverCapabilities(get=False, put=False, cancel=False, probe=False, max_rows_per_op=None) + with pytest.raises(ValueError): + DriverCapabilities(get=True, put=True, cancel=False, probe=False, max_rows_per_op=0) + caps = DriverCapabilities(get=True, put=False, cancel=False, probe=False, max_rows_per_op=4) + with pytest.raises(AttributeError): + caps.put = True # type: ignore[misc] + + +# The in-memory driver's trailer. + + +def test_trailer_round_trip_and_mismatches() -> None: + trailer = encode_trailer(b"name", 64) + assert len(trailer) == TRAILER_BYTES + assert check_trailer(trailer, b"name", 64) is None + assert check_trailer(trailer, b"other", 64) == "object name mismatch" + assert check_trailer(trailer, b"name", 65) == "payload length 64, expected 65" + assert check_trailer(b"\0" * TRAILER_BYTES, b"name", 64) == "bad magic" + assert check_trailer(trailer[:-1], b"name", 64) == ( + f"trailer is {TRAILER_BYTES - 1} bytes, expected {TRAILER_BYTES}" + ) + with pytest.raises(ValueError): + encode_trailer(b"name", 0) + + +def test_trailer_rejects_another_version() -> None: + trailer = bytearray(encode_trailer(b"name", 64)) + trailer[4:6] = (TRAILER_VERSION + 1).to_bytes(2, "little") + assert check_trailer(bytes(trailer), b"name", 64) == ( + f"trailer version {TRAILER_VERSION + 1}, expected {TRAILER_VERSION}" + ) + + +@pytest.mark.parametrize("checksum", list(Checksum)) +def test_every_trailer_byte_is_checked(checksum: Checksum) -> None: + payload = bytes(range(64)) + trailer = encode_trailer(b"name", 64, checksum=checksum, payload=payload) + assert check_trailer(trailer, b"name", 64, payload) is None + for index in range(TRAILER_BYTES): + flipped = bytearray(trailer) + flipped[index] ^= 0x01 + assert check_trailer(bytes(flipped), b"name", 64, payload) is not None, index + + +def test_checksum_covers_the_payload() -> None: + payload = bytes(range(64)) + trailer = encode_trailer(b"name", 64, checksum=Checksum.CRC32, payload=payload) + assert check_trailer(trailer, b"name", 64, payload) is None + corrupted = payload[:-1] + b"\xff" + assert check_trailer(trailer, b"name", 64, corrupted) == "payload checksum mismatch" + assert check_trailer(trailer, b"name", 64) == "checksum needs the payload" + with pytest.raises(ValueError): + encode_trailer(b"name", 64, checksum=Checksum.CRC32) + with pytest.raises(ValueError): + encode_trailer(b"name", 64, checksum=Checksum.CRC32, payload=payload[:-1]) + + +def test_no_checksum_means_a_zero_checksum_field() -> None: + plain = encode_trailer(b"name", 64) + assert check_trailer(plain, b"name", 64, bytes(64)) is None + tampered = plain[:-8] + (1).to_bytes(8, "little") + assert check_trailer(tampered, b"name", 64) == "checksum set without an algorithm" diff --git a/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py new file mode 100644 index 000000000000..c84f5f49a6dc --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py @@ -0,0 +1,1081 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The in-memory driver: the shared conformance checks, its controlled mode, and proof that every +check fails against a driver that breaks exactly the rule it guards.""" + +from __future__ import annotations + +import ctypes +import threading +import time +from collections.abc import Callable, Sequence + +import pytest +from driver_conformance import ( + SLOT, + UNRELEASED_ENDPOINTS, + DriverConformance, + Endpoint, + Harness, + key_of, + statuses, +) +from host_memory import HostBuffer +from memory_store import ( + DEFAULT_CAPABILITIES, + MAX_PROBES_IN_FLIGHT, + Kind, + MemoryDriver, + MemoryOp, + MemoryStore, + StoredObject, +) +from object_trailer import encode_trailer + +from tensorrt_llm._torch.disaggregation.blob.store import ( + MISS, + DriverCapabilities, + DriverRejected, + OpHandle, + ProbeStatus, + RegionHandle, + RegionLayout, + RegistrationError, + RowIO, + RowResult, + RowStatus, + StoreDriver, + StoreError, + concluded, + faulted, +) + +pytestmark = pytest.mark.cpu_only + +KEY_A, KEY_B, KEY_C = key_of("A"), key_of("B"), key_of("C") +TIMEOUT = 5.0 + + +# Stores that misbehave on purpose. + + +class FailingStore(MemoryStore): + """Raises ``StoreError`` on the n-th commit, as a store that drops a connection would.""" + + def __init__(self, fail_on_commit: int) -> None: + super().__init__() + self._fail_on = fail_on_commit + self._commits = 0 + + def commit(self, key: str, payload: bytes, trailer: bytes): # type: ignore[no-untyped-def] + self._commits += 1 + if self._commits == self._fail_on: + raise StoreError("connection dropped") + return super().commit(key, payload, trailer) + + +class TimeoutStore(MemoryStore): + """Reports a timeout for reads of ``key``, as an SDK whose wait expired would.""" + + def __init__(self, key: str) -> None: + super().__init__() + self._key = key + + def get(self, key: str) -> StoredObject | None: + if key == self._key: + raise StoreError("timeout") + return super().get(key) + + +class CrashingStore(MemoryStore): + """Raises something that is not a ``StoreError`` from inside a commit, as a binding bug would.""" + + def commit(self, key: str, payload: bytes, trailer: bytes): # type: ignore[no-untyped-def] + raise RuntimeError("binding returned garbage") + + +class GatedStore(MemoryStore): + """Holds reads and writes of ``held`` keys (every key when None) while the gate is closed, + and says when one is waiting.""" + + def __init__(self, held: set[str] | None = None) -> None: + super().__init__() + self.held = held + self.gate = threading.Event() + self.gate.set() + self.waiting = threading.Event() + + def _hold(self, key: str) -> None: + if (self.held is None or key in self.held) and not self.gate.is_set(): + self.waiting.set() + self.gate.wait() + + def commit(self, key: str, payload: bytes, trailer: bytes): # type: ignore[no-untyped-def] + self._hold(key) + return super().commit(key, payload, trailer) + + def get(self, key: str) -> StoredObject | None: + self._hold(key) + return super().get(key) + + def presence(self, key: str) -> ProbeStatus: + self._hold(key) + return super().presence(key) + + +def _run(target: Callable[..., object], *args: object) -> threading.Thread: + worker = threading.Thread(target=target, args=args, daemon=True) + worker.start() + return worker + + +class TestMemoryDriverConformance(DriverConformance): + driver_cls: type[MemoryDriver] = MemoryDriver + capabilities = DEFAULT_CAPABILITIES + + def make_driver(self) -> StoreDriver: + return self.driver_cls(MemoryStore(), capabilities=self.capabilities) + + def make_blocked_probe(self) -> tuple[StoreDriver, str, Callable[[], None]]: + store = GatedStore(held={KEY_A}) + store.gate.clear() + return self.driver_cls(store, capabilities=self.capabilities), KEY_A, store.gate.set + + +class TestMemoryDriver: + """Behaviour only the in-memory driver's controls or stores can stage.""" + + driver_cls: type[MemoryDriver] = MemoryDriver + + def make(self, store: MemoryStore | None = None, **kwargs: object) -> MemoryDriver: + store = store if store is not None else MemoryStore() + return self.driver_cls(store, **kwargs) # type: ignore[arg-type] + + def _publish( + self, driver: MemoryDriver, rows: dict[str, bytes], *, keep_registered: bool = False + ) -> Endpoint: + writer = Endpoint(driver, list(rows)) + for i, data in enumerate(rows.values()): + writer.fill_row(i, data) + op = driver.put(writer.rows()) + assert statuses(driver, op) == [RowStatus.READY] * len(rows) + if not keep_registered: + writer.unregister() + return writer + + def _shutdown(self, driver: MemoryDriver) -> None: + assert driver.drain(TIMEOUT) is True + for registration in driver.live_registrations(): + driver.unregister(registration) + driver.close() + + # The store and integrity. + + def test_put_keeps_a_private_copy_with_a_trailer(self) -> None: + store = MemoryStore() + self._publish(self.make(store), {"A": b"aaaaaaaa"}) + stored = store.committed_object(KEY_A) + assert stored.payload == b"aaaaaaaa" and stored.trailer + assert len(store) == 1 + + def test_a_corrupted_object_is_a_fault_and_never_lands(self) -> None: + store = MemoryStore() + driver = self.make(store) + self._publish(driver, {"A": b"aaaaaaaa"}) + store.corrupt(KEY_A, 3) + reader = Endpoint(driver, ["A"], fill=0x2D) + (result,) = driver.poll(driver.get(reader.rows())) + assert result.status is RowStatus.FAULT + assert result.fault == "payload checksum mismatch" + assert reader.read_row(0) == b"-" * SLOT + + def test_a_different_object_under_the_same_key_is_diagnosed(self) -> None: + store = MemoryStore() + driver = self.make(store) + self._publish(driver, {"A": b"aaaaaaaa"}) + other = Endpoint(driver, ["A"]) + other.fill_row(0, b"zzzzzzzz") + (result,) = driver.poll(driver.put(other.rows())) + assert result.status is RowStatus.FAULT + assert store.committed_payload(KEY_A) == b"aaaaaaaa" + assert len(store) == 1 + + def test_a_retry_that_fails_in_the_store_leaves_the_committed_object_whole(self) -> None: + store = FailingStore(fail_on_commit=2) + driver = self.make(store) + self._publish(driver, {"A": b"aaaaaaaa"}) + retry = Endpoint(driver, ["A"]) + retry.fill_row(0, b"aaaaaaaa") + (result,) = driver.poll(driver.put(retry.rows())) + assert result.status is RowStatus.FAULT and result.fault == "connection dropped" + reader = Endpoint(driver, ["A"], fill=0x2D) + (got,) = driver.poll(driver.get(reader.rows())) + assert got.status is RowStatus.READY + assert reader.read_row(0) == b"aaaaaaaa" + + def test_readers_never_see_a_partial_object(self) -> None: + store = MemoryStore() + pattern_a, pattern_b = b"a" * 4096, b"b" * 4096 + keys = [f"k{i}" for i in range(200)] + stop = threading.Event() + seen_partial = [] + + def write() -> None: + for i, key in enumerate(keys): + store.commit(key, pattern_a if i % 2 == 0 else pattern_b, b"t") + stop.set() + + def read() -> None: + while not stop.is_set(): + for key in keys: + stored = store.get(key) + if stored is not None and stored.payload not in (pattern_a, pattern_b): + seen_partial.append(key) + + workers = [_run(read), _run(read), _run(write)] + for worker in workers: + worker.join() + assert seen_partial == [] + assert len(store) == len(keys) + + # Statuses that are not faults. + + def test_a_key_held_by_an_unfinished_write_is_not_ready_for_get_put_and_probe(self) -> None: + store = MemoryStore() + driver = self.make(store) + store.begin_write(KEY_A) + reader = Endpoint(driver, ["A", "C"], fill=0x2D) + results = driver.poll(driver.get(reader.rows())) + assert [r.status for r in results] == [RowStatus.NOT_READY, RowStatus.MISS] + assert reader.read_row(0) == b"-" * SLOT + writer = Endpoint(driver, ["A"]) + writer.fill_row(0, b"aaaaaaaa") + assert statuses(driver, driver.put(writer.rows())) == [RowStatus.NOT_READY] + assert driver.probe([KEY_A, KEY_C], TIMEOUT) == (ProbeStatus.NOT_READY, ProbeStatus.MISS) + store.end_write(KEY_A) + assert statuses(driver, driver.put(writer.rows())) == [RowStatus.READY] + assert statuses(driver, driver.get(reader.rows())) == [RowStatus.READY, RowStatus.MISS] + assert reader.read_row(0) == b"aaaaaaaa" + + def test_a_row_the_store_has_no_room_for_is_declined_not_faulted(self) -> None: + driver = self.make(MemoryStore(capacity=1)) + writer = Endpoint(driver, ["A", "B"]) + writer.fill_row(0, b"aaaaaaaa") + writer.fill_row(1, b"bbbbbbbb") + op = driver.put(writer.rows()) + assert statuses(driver, op) == [RowStatus.READY, RowStatus.DECLINED] + assert driver.access_ended(op) + reader = Endpoint(driver, ["A", "B"], fill=0x2D) + assert statuses(driver, driver.get(reader.rows())) == [RowStatus.READY, RowStatus.MISS] + + # Result and access end are independent. + + def test_result_can_come_before_access_end(self) -> None: + store = MemoryStore() + self._publish(self.make(store), {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = self.make(store, controlled=True) + reader = Endpoint(driver, ["A", "B"], fill=0x2D) + op = driver.get(reader.rows()) + assert statuses(driver, op) == [RowStatus.PENDING, RowStatus.PENDING] + + driver.complete(op, [0]) + assert statuses(driver, op) == [RowStatus.READY, RowStatus.PENDING] + assert reader.read_row(0) == b"aaaaaaaa" + assert reader.read_row(1) == b"-" * SLOT + + driver.fail(op, "timeout", [1]) + results = driver.poll(op) + assert concluded(results) + assert results[1] == faulted("timeout") + assert not driver.access_ended(op) + + driver.end_access(op) + assert driver.access_ended(op) + assert driver.poll(op) == results + + def test_complete_accepts_a_one_shot_iterator(self) -> None: + store = MemoryStore() + self._publish(self.make(store), {"A": b"aaaaaaaa"}) + driver = self.make(store, controlled=True) + reader = Endpoint(driver, ["A"], fill=0x2D) + op = driver.get(reader.rows()) + driver.complete(op, iter([0])) + assert statuses(driver, op) == [RowStatus.READY] + assert reader.read_row(0) == b"aaaaaaaa" + + def test_access_can_end_before_the_result_is_published(self) -> None: + store = MemoryStore() + self._publish(self.make(store), {"A": b"aaaaaaaa"}) + driver = self.make(store, controlled=True) + reader = Endpoint(driver, ["A"], fill=0x2D) + op = driver.get(reader.rows()) + driver.transfer(op) + driver.end_access(op) + assert driver.access_ended(op) + assert statuses(driver, op) == [RowStatus.PENDING] + driver.publish_results(op) + assert statuses(driver, op) == [RowStatus.READY] + assert reader.read_row(0) == b"aaaaaaaa" + + def test_a_late_write_lands_only_while_access_is_open(self) -> None: + store = MemoryStore() + self._publish(self.make(store), {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = self.make(store, controlled=True) + reader = Endpoint(driver, ["A", "B"], fill=0x2D) + op = driver.get(reader.rows()) + driver.complete(op, [0]) + driver.fail(op, "timeout", [1]) + results = driver.poll(op) + + driver.write_late(op, 1, b"bbbbbbbb") + assert reader.read_row(1) == b"bbbbbbbb" + assert driver.poll(op) == results + assert not driver.access_ended(op) + + with pytest.raises(RuntimeError, match="ready, not faulted"): + driver.write_late(op, 0, b"cccccccc") + assert reader.read_row(0) == b"aaaaaaaa" + + driver.end_access(op) + with pytest.raises(RuntimeError, match="ended"): + driver.write_late(op, 1, b"cccccccc") + with pytest.raises(RuntimeError, match="ended"): + driver.transfer(op, [1]) + assert reader.read_row(1) == b"bbbbbbbb" + assert reader.intact() + + def test_a_cancelled_row_cannot_be_written_late(self) -> None: + driver = self.make(controlled=True) + reader = Endpoint(driver, ["A"], fill=0x2D) + op = driver.get(reader.rows()) + driver.cancel(op) + assert statuses(driver, op) == [RowStatus.CANCELLED] + with pytest.raises(RuntimeError): + driver.write_late(op, 0, b"aaaaaaaa") + assert reader.read_row(0) == b"-" * SLOT + + # Store errors. + + def test_a_store_timeout_is_a_fault_not_a_miss(self) -> None: + driver = self.make(TimeoutStore(KEY_A)) + reader = Endpoint(driver, ["A", "C"], fill=0x2D) + results = driver.poll(driver.get(reader.rows())) + assert [r.status for r in results] == [RowStatus.FAULT, RowStatus.MISS] + assert results[0].fault == "timeout" + assert reader.read_row(0) == b"-" * SLOT + + def test_a_store_error_mid_operation_becomes_a_row_fault_and_keeps_the_handle(self) -> None: + store = FailingStore(fail_on_commit=2) + driver = self.make(store) + writer = Endpoint(driver, ["A", "B", "C"]) + for i, data in enumerate((b"aaaaaaaa", b"bbbbbbbb", b"cccccccc")): + writer.fill_row(i, data) + op = driver.put(writer.rows()) + results = driver.poll(op) + assert [r.status for r in results] == [RowStatus.READY, RowStatus.FAULT, RowStatus.READY] + assert results[1].fault == "connection dropped" + assert driver.access_ended(op) + assert store.committed_payload(KEY_A) == b"aaaaaaaa" + assert not store.contains(KEY_B) + assert store.committed_payload(KEY_C) == b"cccccccc" + + def test_a_raise_after_the_operation_is_recorded_becomes_a_row_fault(self) -> None: + driver = self.make(CrashingStore()) + writer = Endpoint(driver, ["A"]) + writer.fill_row(0, b"aaaaaaaa") + try: + op = driver.put(writer.rows()) + except Exception as error: # noqa: BLE001 - the check is that nothing escapes + pytest.fail(f"put raised after the operation was recorded: {error!r}") + (result,) = driver.poll(op) + assert result.status is RowStatus.FAULT + assert result.fault == "RuntimeError: binding returned garbage" + assert driver.access_ended(op) + assert driver.drain(TIMEOUT) is True + + # Concurrency. + + def test_queries_do_not_wait_on_an_operation_inside_the_store(self) -> None: + store = GatedStore() + driver = self.make(store, controlled=True) + done = Endpoint(driver, ["A"]) + finished = driver.get(done.rows()) + driver.fail(finished, "timeout") + + slow = Endpoint(driver, ["B"]) + slow.fill_row(0, b"bbbbbbbb") + blocked = driver.put(slow.rows()) + store.gate.clear() + worker = _run(driver.transfer, blocked) + answered = threading.Event() + seen: dict[str, object] = {} + failure: list[Exception] = [] + + def query() -> None: + # Off the main thread: if a query waits on the store, the main thread can still open + # the gate and the test fails instead of hanging. + try: + seen["finished"] = statuses(driver, finished) + seen["blocked"] = statuses(driver, blocked) + seen["access_ended"] = driver.access_ended(blocked) + try: + driver.end_access(blocked) + except RuntimeError as error: + seen["end_access"] = str(error) + driver.cancel(blocked) + seen["after_cancel"] = statuses(driver, blocked) + except Exception as error: # surfaced on the main thread below + failure.append(error) + finally: + answered.set() + + querier = None + try: + assert store.waiting.wait(TIMEOUT), "the worker never reached the store" + querier = _run(query) + assert answered.wait(TIMEOUT), "queries waited on an operation inside the store" + finally: + store.gate.set() + worker.join(TIMEOUT) + if querier is not None: + querier.join(TIMEOUT) + assert not worker.is_alive() + if failure: + raise AssertionError("a query raised") from failure[0] + assert seen["finished"] == [RowStatus.FAULT] + assert seen["blocked"] == [RowStatus.PENDING] + assert seen["access_ended"] is False + assert seen.get("end_access") == "a transfer is still in progress" + assert seen["after_cancel"] == [RowStatus.PENDING] + driver.publish_results(blocked) + assert statuses(driver, blocked) == [RowStatus.READY] + assert store.committed_payload(KEY_B) == b"bbbbbbbb" + + def test_a_blocked_get_does_not_delay_another_get(self) -> None: + store = GatedStore(held={KEY_A}) + driver = self.make(store) + self._publish(driver, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + slow = Endpoint(driver, ["A"], fill=0x2D) + quick = Endpoint(driver, ["B"], fill=0x2D) + store.gate.clear() + outcome: dict[str, list[RowStatus]] = {} + + def fetch(name: str, endpoint: Endpoint) -> None: + outcome[name] = statuses(driver, driver.get(endpoint.rows())) + + slow_worker = _run(fetch, "slow", slow) + try: + assert store.waiting.wait(TIMEOUT), "the slow get never reached the store" + quick_worker = _run(fetch, "quick", quick) + quick_worker.join(TIMEOUT) + assert not quick_worker.is_alive(), "a get waited for another get blocked in the store" + assert outcome["quick"] == [RowStatus.READY] + assert quick.read_row(0) == b"bbbbbbbb" + assert "slow" not in outcome + finally: + store.gate.set() + slow_worker.join(TIMEOUT) + assert outcome["slow"] == [RowStatus.READY] + + # Cancel. + + def test_cancel_concludes_unstarted_rows_and_spares_started_ones(self) -> None: + store = MemoryStore() + self._publish(self.make(store), {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = self.make(store, controlled=True) + reader = Endpoint(driver, ["A", "B"], fill=0x2D) + op = driver.get(reader.rows()) + driver.transfer(op, [0]) + driver.cancel(op) + assert statuses(driver, op) == [RowStatus.PENDING, RowStatus.CANCELLED] + driver.publish_results(op, [0]) + assert statuses(driver, op) == [RowStatus.READY, RowStatus.CANCELLED] + assert reader.read_row(1) == b"-" * SLOT + assert not driver.access_ended(op) + + def test_cancelling_everything_before_it_started_ends_access(self) -> None: + driver = self.make(controlled=True) + reader = Endpoint(driver, ["A", "B"]) + op = driver.get(reader.rows()) + driver.cancel(op) + assert statuses(driver, op) == [RowStatus.CANCELLED, RowStatus.CANCELLED] + assert driver.access_ended(op) + + def test_an_ineffective_cancel_changes_nothing(self) -> None: + driver = self.make(controlled=True) + driver.set_cancel_effective(False) + reader = Endpoint(driver, ["A"]) + op = driver.get(reader.rows()) + driver.cancel(op) + assert statuses(driver, op) == [RowStatus.PENDING] + assert not driver.access_ended(op) + + # Probe. + + def test_probe_answers_unknown_when_the_store_does_not_answer_in_time(self) -> None: + store = GatedStore(held={KEY_A}) + driver = self.make(store) + self._publish(driver, {"B": b"bbbbbbbb"}) + keys = [KEY_B, KEY_A, KEY_C] + store.gate.clear() + answer: list[tuple[ProbeStatus, ...]] = [] + elapsed: list[float] = [] + + def probe() -> None: + started = time.monotonic() + answer.append(driver.probe(keys, 0.05)) + elapsed.append(time.monotonic() - started) + + prober = _run(probe) + prober.join(1.0) + try: + assert not prober.is_alive(), "probe waited for the store past its budget" + # Keys asked before the blocked one were answered; the rest are unknown, not MISS. + assert answer == [(ProbeStatus.HIT, ProbeStatus.UNKNOWN, ProbeStatus.UNKNOWN)] + assert elapsed[0] < 0.5 + assert driver.probes_in_flight() == 1 + assert driver.drain(0.05) is False # the probe is still inside the store + finally: + store.gate.set() + prober.join(TIMEOUT) + driver.join_probes(TIMEOUT) + assert driver.probe(keys, TIMEOUT) == (ProbeStatus.UNKNOWN,) * 3 # draining now + self._shutdown(driver) + + def test_probes_in_flight_are_capped(self) -> None: + store = GatedStore(held={KEY_A}) + driver = self.make(store) + store.gate.clear() + try: + for _ in range(MAX_PROBES_IN_FLIGHT): + assert driver.probe([KEY_A], 0.0) == (ProbeStatus.UNKNOWN,) + assert driver.probes_in_flight() == MAX_PROBES_IN_FLIGHT + started = time.monotonic() + assert driver.probe([KEY_A], 0.2) == (ProbeStatus.UNKNOWN,) + assert time.monotonic() - started < 0.1 # answered at once, no new worker + assert driver.probes_in_flight() == MAX_PROBES_IN_FLIGHT + finally: + store.gate.set() + driver.join_probes(TIMEOUT) + self._shutdown(driver) + assert driver.probes_in_flight() == 0 + + def test_probe_unknown_switch(self) -> None: + driver = self.make() + self._publish(driver, {"A": b"aaaaaaaa"}) + driver.set_probe_unknown(True) + assert driver.probe([KEY_A, KEY_C], TIMEOUT) == (ProbeStatus.UNKNOWN,) * 2 + + def test_a_drain_racing_a_submission_leaves_a_cancelled_operation_with_its_handle( + self, + ) -> None: + store = MemoryStore() + racing = type( + "Racing", (self.driver_cls,), {"_before_execute": lambda self, op: self.drain(0.0)} + ) + driver = racing(store) + writer = Endpoint(driver, ["A"]) + writer.fill_row(0, b"aaaaaaaa") + try: + op = driver.put(writer.rows()) + except Exception as error: # noqa: BLE001 - the check is that the handle comes back + pytest.fail(f"put raised instead of returning its handle: {error!r}") + assert statuses(driver, op) == [RowStatus.CANCELLED] + assert driver.access_ended(op) + assert len(store) == 0 + assert driver.drain(TIMEOUT) is True + + def test_an_object_shorter_than_its_trailer_claims_is_a_fault_and_never_lands(self) -> None: + store = MemoryStore() + driver = self.make(store) + store.overwrite(KEY_A, b"aaaa", encode_trailer(KEY_A.encode(), SLOT)) # claims 8 bytes + reader = Endpoint(driver, ["A"], fill=0x2D) + (result,) = driver.poll(driver.get(reader.rows())) + assert result.status is RowStatus.FAULT + assert result.fault == f"object is 4 bytes, expected {SLOT}" + assert reader.read_row(0) == b"-" * SLOT + + # The conformance harness. + + def test_harness_keeps_a_buffer_whose_registration_escaped(self) -> None: + driver = self.make() + h = Harness(driver, TIMEOUT) + driver.fail_next_register("timed out after the request", escaped=True) + with pytest.raises(RegistrationError): + h.endpoint(["A"]) + assert len(h.endpoints) == 1 and h.endpoints[0].handle is not None + assert h.cleanup() == [] + assert driver.live_registrations() == [] + + def test_harness_retains_an_endpoint_it_could_not_unregister_and_retries(self) -> None: + driver = self.make() + h = Harness(driver, TIMEOUT) + endpoint = h.endpoint(["A"]) + assert endpoint.handle is not None + driver.fail_next_unregister(endpoint.handle, "busy") + problems = h.cleanup() + assert problems and not h.done + assert endpoint in UNRELEASED_ENDPOINTS + assert driver.live_registrations() == [endpoint.handle] + assert h.cleanup() == [] and h.done + assert driver.live_registrations() == [] + + # Registration and shutdown. + + def test_unregister_is_refused_while_an_operation_may_access_the_region(self) -> None: + driver = self.make(controlled=True) + reader = Endpoint(driver, ["A"]) + other = Endpoint(driver, ["B"]) + op = driver.get(reader.rows()) + with pytest.raises(DriverRejected): + reader.unregister() + assert reader.handle in driver.live_registrations() + other.unregister() # another region is not held by this operation + driver.transfer(op) + driver.publish_results(op) + with pytest.raises(DriverRejected): + reader.unregister() # a result is not access-end evidence + driver.end_access(op) + reader.unregister() + assert driver.live_registrations() == [] + + def test_drain_waits_for_a_started_transfer(self) -> None: + store = GatedStore(held={KEY_A}) + driver = self.make(store) + self._publish(driver, {"A": b"aaaaaaaa"}) + reader = Endpoint(driver, ["A"], fill=0x2D) + store.gate.clear() + worker = _run(driver.get, reader.rows()) + try: + assert store.waiting.wait(TIMEOUT), "the get never reached the store" + started = time.monotonic() + assert driver.drain(0.05) is False + assert time.monotonic() - started < 1.0 + with pytest.raises(DriverRejected, match="draining"): + driver.get(reader.rows()) + finally: + store.gate.set() + worker.join(TIMEOUT) + assert driver.drain(TIMEOUT) is True + assert reader.read_row(0) == b"aaaaaaaa" + + def test_an_unprovable_operation_lets_drain_return_false_at_once(self) -> None: + driver = self.make(controlled=True) + reader = Endpoint(driver, ["A"]) + op = driver.get(reader.rows()) + driver.transfer(op) + driver.publish_results(op) + driver.mark_unprovable(op) + assert driver.access_unprovable(op) + assert not driver.access_ended(op) + + started = time.monotonic() + assert driver.drain(TIMEOUT) is False + assert time.monotonic() - started < 1.0 + with pytest.raises(DriverRejected): + reader.unregister() + with pytest.raises(DriverRejected): + driver.close() + + driver.end_access(op) # later evidence, for example from the store's own teardown + assert driver.drain(TIMEOUT) is True + reader.unregister() + driver.close() + + def test_drain_returns_false_at_once_even_while_another_operation_is_still_provable( + self, + ) -> None: + driver = self.make(controlled=True) + lost = Endpoint(driver, ["A"]) + pending = Endpoint(driver, ["B"]) + unprovable = driver.get(lost.rows()) + driver.transfer(unprovable) + driver.publish_results(unprovable) + driver.mark_unprovable(unprovable) + provable = driver.get(pending.rows()) + driver.transfer(provable) # started, so drain cannot cancel it; its end is still to come + + started = time.monotonic() + assert driver.drain(1.0) is False + assert time.monotonic() - started < 0.5 + driver.end_access(provable) + assert driver.drain(0.05) is False # the unprovable one still stands + driver.end_access(unprovable) + assert driver.drain(TIMEOUT) is True + + def test_register_checks_geometry_and_overlap(self) -> None: + driver = self.make() + buffer = HostBuffer(64) + with pytest.raises(DriverRejected, match="layout needs"): + driver.register(buffer.address, 64, RegionLayout(SLOT, 9)) + with pytest.raises(DriverRejected): + driver.register(buffer.address, 0, RegionLayout(SLOT, 1)) + handle = driver.register(buffer.address, 64, RegionLayout(SLOT, 8)) + with pytest.raises(DriverRejected, match="overlaps"): + driver.register(buffer.address + 8, 8, RegionLayout(8, 1)) + driver.unregister(handle) + assert driver.live_registrations() == [] + + def test_a_registration_the_store_may_hold_keeps_its_handle(self) -> None: + driver = self.make() + buffer = HostBuffer(64) + driver.fail_next_register("store unreachable", escaped=False) + with pytest.raises(RegistrationError) as clean: + driver.register(buffer.address, 64, RegionLayout(SLOT, 8)) + assert clean.value.handle is None + assert driver.live_registrations() == [] + + driver.fail_next_register("timed out after the request", escaped=True) + with pytest.raises(RegistrationError) as ambiguous: + driver.register(buffer.address, 64, RegionLayout(SLOT, 8)) + handle = ambiguous.value.handle + assert handle is not None + assert driver.live_registrations() == [handle] + + driver.fail_next_unregister(handle, "busy") + with pytest.raises(StoreError, match="busy"): + driver.unregister(handle) + assert driver.live_registrations() == [handle] + driver.unregister(handle) + assert driver.live_registrations() == [] + + def test_calls_after_close_are_programming_errors(self) -> None: + driver = self.make() + writer = self._publish(driver, {"A": b"aaaaaaaa"}, keep_registered=True) + op = driver.get(writer.rows()) + assert driver.drain(TIMEOUT) is True + writer.unregister() + driver.close() + with pytest.raises(RuntimeError, match="closed"): + driver.poll(op) + with pytest.raises(RuntimeError, match="closed"): + driver.drain(0.0) + with pytest.raises(RuntimeError, match="closed"): + driver.probe([KEY_A], 0.0) + with pytest.raises(DriverRejected, match="closed"): + driver.register(writer.payload.address, SLOT, RegionLayout(SLOT, 1)) + + def test_memory_driver_is_a_store_driver(self) -> None: + driver: StoreDriver = self.make() + assert isinstance(driver, StoreDriver) + + +# Drivers that break exactly one rule. Each proves that the check guarding the rule fails +# against it, so a weakened check is caught here before a real driver meets it. + + +class _AcceptsDuplicateKeys(MemoryDriver): + def _check_unique_keys(self, rows: Sequence[RowIO]) -> None: + pass + + +class _TouchesMemoryBeforeChecking(MemoryDriver): + def get(self, rows: Sequence[RowIO]) -> MemoryOp: + for io in rows: + ctypes.memset(io.payload.address, 0, io.payload.nbytes) + return super().get(rows) + + +class _MissTouchesTheDestination(MemoryDriver): + def _get_row(self, io: RowIO) -> RowResult: + result = super()._get_row(io) + if result.status is RowStatus.MISS: + ctypes.memset(io.payload.address, 0, io.payload.nbytes) + return result + + +class _CallsAWrongSizeAMiss(MemoryDriver): + def _get_row(self, io: RowIO) -> RowResult: + result = super()._get_row(io) + return MISS if result.status is RowStatus.FAULT else result + + +class _OverwritesExistingObjects(MemoryDriver): + def _put_row(self, io: RowIO) -> RowResult: + payload = ctypes.string_at(io.payload.address, io.payload.nbytes) + self.store.overwrite(io.key, payload, self._trailer_for(io, payload)) + return super()._put_row(io) + + +class _IgnoresTheRowLimit(MemoryDriver): + def _check_row_limit(self, rows: Sequence[RowIO]) -> None: + if not rows: + raise DriverRejected("an operation needs at least one row") + + +class _IgnoresTheDirection(MemoryDriver): + def _check_direction(self, kind: Kind) -> None: + pass + + +class _NeverEndsAccess(MemoryDriver): + def _execute(self, op: MemoryOp) -> None: + super()._execute(op) + with self._lock: + op.access_ended = False # the evidence is lost again + + +class _ChangesItsCapabilities(MemoryDriver): + _flip = False + + @property + def capabilities(self) -> DriverCapabilities: + self._flip = not self._flip + return DriverCapabilities( + get=True, put=True, cancel=self._flip, probe=True, max_rows_per_op=None + ) + + +class _ProbeTakesItsTime(MemoryDriver): + def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...]: + time.sleep(budget_s + 0.8) + return super().probe(keys, budget_s) + + +class _KeepsAcceptingWhileDraining(MemoryDriver): + def _check_not_draining(self) -> None: + pass + + +class _ClosesAnyTime(MemoryDriver): + def _check_closable(self) -> None: + pass + + +class _UnregisterKeepsTheRange(MemoryDriver): + def unregister(self, handle: RegionHandle) -> None: + pass + + +class _UnregistersBusyRegions(MemoryDriver): + def _check_region_idle(self, registration: object) -> None: + pass + + +class _CallsATimeoutAMiss(MemoryDriver): + def _move(self, kind: Kind, io: RowIO) -> RowResult: + try: + return super()._move(kind, io) + finally: + pass + + def _get_row(self, io: RowIO) -> RowResult: + try: + return super()._get_row(io) + except StoreError: + return MISS + + +class _LetsRaisesEscape(MemoryDriver): + def _move(self, kind: Kind, io: RowIO) -> RowResult: + try: + return self._put_row(io) if kind == "put" else self._get_row(io) + except StoreError as error: + return faulted(str(error)) + + +class _SkipsTheIntegrityCheck(MemoryDriver): + def _verify(self, io: RowIO, stored: StoredObject) -> str | None: + return None + + +class _CallsNotReadyAFault(MemoryDriver): + def _get_row(self, io: RowIO) -> RowResult: + result = super()._get_row(io) + return faulted("not ready") if result.status is RowStatus.NOT_READY else result + + +class _CallsDeclinedAFault(MemoryDriver): + def _put_row(self, io: RowIO) -> RowResult: + result = super()._put_row(io) + return faulted("no space") if result.status is RowStatus.DECLINED else result + + +class _ProbeIgnoresTheBudget(MemoryDriver): + def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...]: + return tuple(self.store.presence(key) for key in keys) + + +class _ProbesWithoutLimit(MemoryDriver): + def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...]: + with self._lock: + self._probes = self._probes[: MAX_PROBES_IN_FLIGHT - 1] + return super().probe(keys, budget_s) + + +class _ReportsAccessEndedEarly(MemoryDriver): + def access_ended(self, op: OpHandle) -> bool: + return True + + +class _ExecutesWithTheControls(MemoryDriver): + def _execute(self, op: MemoryOp) -> None: + self.transfer(op) + self.publish_results(op) + self.end_access(op) + + +class _DrainWaitsForEveryone(MemoryDriver): + def _drain_can_return(self) -> bool: + return self._quiet() + + +LIMITED = DriverCapabilities(get=True, put=True, cancel=True, probe=True, max_rows_per_op=1) +GET_ONLY = DriverCapabilities(get=True, put=False, cancel=True, probe=True, max_rows_per_op=None) + +# rule -> (driver breaking it, the capabilities to run with, the conformance check that must fail) +BROKEN_CONFORMANCE = { + "a key appears once per operation": ( + _AcceptsDuplicateKeys, + DEFAULT_CAPABILITIES, + "test_a_key_appears_at_most_once_per_operation", + ), + "rejection happens before any effect": ( + _TouchesMemoryBeforeChecking, + DEFAULT_CAPABILITIES, + "test_rows_outside_a_live_registration_are_rejected_before_any_effect", + ), + "a miss leaves the destination alone": ( + _MissTouchesTheDestination, + DEFAULT_CAPABILITIES, + "test_get_reports_hit_hit_miss_and_leaves_the_miss_untouched", + ), + "a wrong size is a fault": ( + _CallsAWrongSizeAMiss, + DEFAULT_CAPABILITIES, + "test_an_object_of_the_wrong_size_is_a_fault_not_a_miss", + ), + "a put never damages an existing object": ( + _OverwritesExistingObjects, + DEFAULT_CAPABILITIES, + "test_a_put_under_an_existing_key_leaves_the_object_whole", + ), + "the row limit is enforced": ( + _IgnoresTheRowLimit, + LIMITED, + "test_rows_beyond_the_declared_limit_are_rejected", + ), + "an undeclared direction is rejected": ( + _IgnoresTheDirection, + GET_ONLY, + "test_an_undeclared_direction_is_rejected", + ), + "access ends with evidence": ( + _NeverEndsAccess, + DEFAULT_CAPABILITIES, + "test_results_are_stable_and_access_ends_with_evidence", + ), + "capabilities are fixed": ( + _ChangesItsCapabilities, + DEFAULT_CAPABILITIES, + "test_capabilities_do_not_change", + ), + "probe returns within its budget": ( + _ProbeTakesItsTime, + DEFAULT_CAPABILITIES, + "test_probe_returns_within_its_budget", + ), + "probe answers unknown for a blocked key": ( + _ProbeIgnoresTheBudget, + DEFAULT_CAPABILITIES, + "test_probe_answers_unknown_for_a_key_the_store_does_not_answer_in_time", + ), + "drain is one-way": ( + _KeepsAcceptingWhileDraining, + DEFAULT_CAPABILITIES, + "test_drain_is_one_way", + ), + "close needs a true drain": ( + _ClosesAnyTime, + DEFAULT_CAPABILITIES, + "test_close_needs_a_true_drain_and_no_live_registrations", + ), + "unregister frees the range": ( + _UnregisterKeepsTheRange, + DEFAULT_CAPABILITIES, + "test_unregister_is_idempotent_and_frees_the_range", + ), +} + +# rule -> (driver breaking it, the memory-driver check that must fail) +BROKEN_MEMORY = { + "unregister needs an idle region": ( + _UnregistersBusyRegions, + "test_unregister_is_refused_while_an_operation_may_access_the_region", + ), + "a store timeout is a fault": ( + _CallsATimeoutAMiss, + "test_a_store_timeout_is_a_fault_not_a_miss", + ), + "a recorded operation never raises": ( + _LetsRaisesEscape, + "test_a_raise_after_the_operation_is_recorded_becomes_a_row_fault", + ), + "integrity is verified on every get": ( + _SkipsTheIntegrityCheck, + "test_a_corrupted_object_is_a_fault_and_never_lands", + ), + "an unfinished write is not ready": ( + _CallsNotReadyAFault, + "test_a_key_held_by_an_unfinished_write_is_not_ready_for_get_put_and_probe", + ), + "no room is declined": ( + _CallsDeclinedAFault, + "test_a_row_the_store_has_no_room_for_is_declined_not_faulted", + ), + "a submission racing a drain keeps its handle": ( + _ExecutesWithTheControls, + "test_a_drain_racing_a_submission_leaves_a_cancelled_operation_with_its_handle", + ), + "probes in flight are capped": ( + _ProbesWithoutLimit, + "test_probes_in_flight_are_capped", + ), + "access does not end early": ( + _ReportsAccessEndedEarly, + "test_queries_do_not_wait_on_an_operation_inside_the_store", + ), + "drain returns at once on an unprovable operation": ( + _DrainWaitsForEveryone, + "test_drain_returns_false_at_once_even_while_another_operation_is_still_provable", + ), +} + +CHECK_FAILED = (AssertionError, pytest.fail.Exception) + + +def _conformance(driver_cls: type[MemoryDriver], caps: DriverCapabilities) -> DriverConformance: + suite = type( + "Suite", + (TestMemoryDriverConformance,), + {"driver_cls": driver_cls, "capabilities": caps, "timeout": 0.3}, + ) + return suite() # type: ignore[no-any-return] + + +@pytest.mark.parametrize("rule", sorted(BROKEN_CONFORMANCE)) +def test_each_conformance_check_passes_the_driver_and_fails_the_one_breaking_its_rule( + rule: str, +) -> None: + broken, capabilities, check = BROKEN_CONFORMANCE[rule] + getattr(_conformance(MemoryDriver, capabilities), check)() + with pytest.raises(CHECK_FAILED): + getattr(_conformance(broken, capabilities), check)() + + +@pytest.mark.parametrize("rule", sorted(BROKEN_MEMORY)) +def test_each_memory_driver_check_fails_the_driver_breaking_its_rule(rule: str) -> None: + broken, check = BROKEN_MEMORY[rule] + suite = type("Suite", (TestMemoryDriver,), {"driver_cls": broken})() + with pytest.raises(CHECK_FAILED): + getattr(suite, check)() + + +def test_each_broken_driver_breaks_one_rule_and_names_an_existing_check() -> None: + conformance = [cls for cls, _, _ in BROKEN_CONFORMANCE.values()] + memory = [cls for cls, _ in BROKEN_MEMORY.values()] + assert len(set(conformance + memory)) == len(conformance) + len(memory) + checks = {check for _, _, check in BROKEN_CONFORMANCE.values()} + assert checks <= {name for name in dir(DriverConformance) if name.startswith("test_")} + checks = {check for _, check in BROKEN_MEMORY.values()} + assert checks <= {name for name in dir(TestMemoryDriver) if name.startswith("test_")} diff --git a/tests/unittest/_torch/disaggregation/blob/test_memory_store.py b/tests/unittest/_torch/disaggregation/blob/test_memory_store.py deleted file mode 100644 index 237024e40751..000000000000 --- a/tests/unittest/_torch/disaggregation/blob/test_memory_store.py +++ /dev/null @@ -1,566 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The private driver boundary, exercised through the in-memory store and driver. - -These tests pin what any driver owes the backend: bytes move whole and never alias the source, -results are per row and never pass a timeout off as a miss, access end is separate evidence, and -failed registrations stay cleanable. -""" - -from __future__ import annotations - -import threading - -import pytest -from host_memory import HostBuffer -from memory_store import MemoryDriver, MemoryOp, MemoryStore - -from tensorrt_llm._torch.disaggregation.blob.format import ( - TRAILER_BYTES, - check_trailer, - encode_trailer, -) -from tensorrt_llm._torch.disaggregation.blob.keys import KEY_FORMAT_VERSION, encode_key -from tensorrt_llm._torch.disaggregation.blob.store import ( - ProbeStatus, - RegionLayout, - RegistrationError, - RowIO, - RowStatus, - Span, - StoreDriver, - StoreError, - concluded, -) - -pytestmark = pytest.mark.cpu_only - -SLOT = 8 -NAMES = {"A": b"name-A", "B": b"name-B", "C": b"name-C"} - - -class Endpoint: - """One side of a transfer: a registered payload region plus registered trailer buffers.""" - - def __init__(self, driver: MemoryDriver, labels: list[str], fill: int = 0) -> None: - self.driver = driver - self.labels = labels - self.payload = HostBuffer(SLOT * len(labels), fill) - self.trailer = HostBuffer(TRAILER_BYTES * len(labels), fill) - self.handles = [ - driver.register( - self.payload.address, self.payload.nbytes, RegionLayout(SLOT, len(labels), "kv") - ), - driver.register( - self.trailer.address, - self.trailer.nbytes, - RegionLayout(TRAILER_BYTES, len(labels), "trailer"), - ), - ] - self.rows = [ - RowIO( - encode_key("t", NAMES[label]), - self.payload.span(i * SLOT, SLOT), - self.trailer.span(i * TRAILER_BYTES, TRAILER_BYTES), - (0, i), - ) - for i, label in enumerate(labels) - ] - - def fill_row(self, index: int, data: bytes) -> None: - self.payload.write(data, index * SLOT) - self.trailer.write(encode_trailer(NAMES[self.labels[index]], SLOT), index * TRAILER_BYTES) - - def read_row(self, index: int) -> bytes: - return self.payload.read(index * SLOT, SLOT) - - def trailer_ok(self, index: int) -> str | None: - trailer = self.trailer.read(index * TRAILER_BYTES, TRAILER_BYTES) - return check_trailer(trailer, NAMES[self.labels[index]], SLOT) - - def intact(self) -> bool: - return self.payload.canaries_intact() and self.trailer.canaries_intact() - - -def publish(store: MemoryStore, labels: dict[str, bytes]) -> None: - """Put ``labels`` into ``store`` through an automatic driver of its own.""" - writer = Endpoint(MemoryDriver(store), list(labels)) - for i, data in enumerate(labels.values()): - writer.fill_row(i, data) - op = writer.driver.put(writer.rows) - assert all(r.status is RowStatus.READY for r in writer.driver.poll(op)) - - -@pytest.fixture -def store() -> MemoryStore: - return MemoryStore() - - -def statuses(driver: MemoryDriver, op: MemoryOp) -> list[RowStatus]: - return [r.status for r in driver.poll(op)] - - -# Bytes and ownership. - - -def test_put_keeps_its_own_copy_so_the_source_may_change(store: MemoryStore) -> None: - writer = Endpoint(MemoryDriver(store), ["A"]) - writer.fill_row(0, b"aaaaaaaa") - op = writer.driver.put(writer.rows) - assert statuses(writer.driver, op) == [RowStatus.READY] - assert writer.driver.access_ended(op) - - writer.payload.write(b"xxxxxxxx") - assert store.committed_payload(writer.rows[0].key) == b"aaaaaaaa" - - reader = Endpoint(MemoryDriver(store), ["A"], fill=0x2D) - got = reader.driver.get(reader.rows) - assert statuses(reader.driver, got) == [RowStatus.READY] - assert reader.read_row(0) == b"aaaaaaaa" - assert reader.trailer_ok(0) is None - assert reader.intact() - - -def test_get_reports_hit_hit_miss_per_row_and_leaves_the_miss_untouched(store: MemoryStore) -> None: - publish(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) - reader = Endpoint(MemoryDriver(store), ["A", "B", "C"], fill=0x2D) - op = reader.driver.get(reader.rows) - assert statuses(reader.driver, op) == [RowStatus.READY, RowStatus.READY, RowStatus.MISS] - assert concluded(reader.driver.poll(op)) - assert reader.read_row(0) == b"aaaaaaaa" - assert reader.read_row(1) == b"bbbbbbbb" - assert reader.read_row(2) == b"-" * SLOT - assert reader.trailer_ok(2) is not None - assert reader.intact() - - -def test_get_rejects_a_size_mismatch_before_writing(store: MemoryStore) -> None: - key = encode_key("t", NAMES["A"]) - assert store.commit(key, b"short") - reader = Endpoint(MemoryDriver(store), ["A"], fill=0x2D) - op = reader.driver.get(reader.rows) - (result,) = reader.driver.poll(op) - assert result.status is RowStatus.FAULT - assert result.fault == f"object is 5 bytes, expected {SLOT + TRAILER_BYTES}" - assert reader.read_row(0) == b"-" * SLOT - assert reader.intact() - - -def test_rows_outside_a_live_registration_are_refused_before_any_effect(store: MemoryStore) -> None: - driver = MemoryDriver(store) - ok = Endpoint(driver, ["A"]) - stray = HostBuffer(SLOT) - row = RowIO(ok.rows[0].key, stray.span(), ok.rows[0].trailer) - with pytest.raises(ValueError, match="outside"): - driver.get([row]) - assert driver.open_ops() == [] - driver.unregister(ok.handles[0]) - with pytest.raises(ValueError, match="outside"): - driver.get(ok.rows) - - -# Format. - - -def test_trailer_round_trip_and_mismatches() -> None: - trailer = encode_trailer(b"name", 64) - assert len(trailer) == TRAILER_BYTES - assert check_trailer(trailer, b"name", 64) is None - assert check_trailer(trailer, b"other", 64) == "object name mismatch" - assert check_trailer(trailer, b"name", 65) == "payload length 64, expected 65" - assert check_trailer(b"\0" * TRAILER_BYTES, b"name", 64) == "bad magic" - assert check_trailer(trailer[:-1], b"name", 64) == ( - f"trailer is {TRAILER_BYTES - 1} bytes, expected {TRAILER_BYTES}" - ) - - -def test_key_encoding_is_prefix_version_and_hex() -> None: - key = encode_key("prod", b"\x01\xff") - assert key == f"prod/v{KEY_FORMAT_VERSION}/01ff" - for prefix in ("", "a/b", "a b"): - with pytest.raises(ValueError): - encode_key(prefix, b"\x01") - with pytest.raises(ValueError): - encode_key("prod", b"") - - -def test_row_refuses_overlapping_payload_and_trailer() -> None: - with pytest.raises(ValueError, match="overlap"): - RowIO("k", Span(0x1000, 64), Span(0x1020, 32)) - - -# Publication. - - -def test_duplicate_put_reuses_identical_object_and_rejects_different_bytes( - store: MemoryStore, -) -> None: - first = Endpoint(MemoryDriver(store), ["A"]) - first.fill_row(0, b"aaaaaaaa") - assert statuses(first.driver, first.driver.put(first.rows)) == [RowStatus.READY] - - same = Endpoint(MemoryDriver(store), ["A"]) - same.fill_row(0, b"aaaaaaaa") - assert statuses(same.driver, same.driver.put(same.rows)) == [RowStatus.READY] - - other = Endpoint(MemoryDriver(store), ["A"]) - other.fill_row(0, b"zzzzzzzz") - (result,) = other.driver.poll(other.driver.put(other.rows)) - assert result.status is RowStatus.FAULT - assert store.committed_payload(first.rows[0].key) == b"aaaaaaaa" - assert len(store) == 1 - - -def test_a_failed_retry_does_not_damage_the_committed_object(store: MemoryStore) -> None: - publish(store, {"A": b"aaaaaaaa"}) - driver = MemoryDriver(store, controlled=True) - retry = Endpoint(driver, ["A"]) - retry.fill_row(0, b"aaaaaaaa") - op = driver.put(retry.rows) - driver.fail(op, "connection reset") - (result,) = driver.poll(op) - assert result.status is RowStatus.FAULT - assert not driver.access_ended(op) - assert store.committed_payload(retry.rows[0].key) == b"aaaaaaaa" - - -def test_readers_never_see_a_partial_object(store: MemoryStore) -> None: - pattern_a, pattern_b = b"a" * 4096, b"b" * 4096 - keys = [f"k{i}" for i in range(200)] - stop = threading.Event() - seen_partial = [] - - def write() -> None: - for i, key in enumerate(keys): - store.commit(key, pattern_a if i % 2 == 0 else pattern_b) - stop.set() - - def read() -> None: - while not stop.is_set(): - for key in keys: - data = store.get(key) - if data is not None and data not in (pattern_a, pattern_b): - seen_partial.append(key) - - readers = [threading.Thread(target=read) for _ in range(2)] - writer = threading.Thread(target=write) - for t in readers + [writer]: - t.start() - for t in readers + [writer]: - t.join() - assert seen_partial == [] - assert len(store) == len(keys) - - -# Result and access end are independent. - - -def test_result_can_come_before_access_end(store: MemoryStore) -> None: - publish(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) - driver = MemoryDriver(store, controlled=True) - reader = Endpoint(driver, ["A", "B"], fill=0x2D) - op = driver.get(reader.rows) - assert statuses(driver, op) == [RowStatus.PENDING, RowStatus.PENDING] - - driver.complete(op, [0]) - assert statuses(driver, op) == [RowStatus.READY, RowStatus.PENDING] - assert reader.read_row(0) == b"aaaaaaaa" - assert reader.read_row(1) == b"-" * SLOT - - driver.fail(op, "timeout", [1]) - results = driver.poll(op) - assert concluded(results) - assert results[1].status is RowStatus.FAULT and results[1].fault == "timeout" - assert not driver.access_ended(op) - - driver.end_access(op) - assert driver.access_ended(op) - assert driver.poll(op) == results - - -def test_complete_accepts_a_one_shot_iterator(store: MemoryStore) -> None: - publish(store, {"A": b"aaaaaaaa"}) - driver = MemoryDriver(store, controlled=True) - reader = Endpoint(driver, ["A"], fill=0x2D) - op = driver.get(reader.rows) - driver.complete(op, iter([0])) - assert statuses(driver, op) == [RowStatus.READY] - assert reader.read_row(0) == b"aaaaaaaa" - - -def test_access_can_end_before_the_result_is_published(store: MemoryStore) -> None: - publish(store, {"A": b"aaaaaaaa"}) - driver = MemoryDriver(store, controlled=True) - reader = Endpoint(driver, ["A"], fill=0x2D) - op = driver.get(reader.rows) - driver.transfer(op) - driver.end_access(op) - assert driver.access_ended(op) - assert statuses(driver, op) == [RowStatus.PENDING] - driver.publish_results(op) - assert statuses(driver, op) == [RowStatus.READY] - assert reader.read_row(0) == b"aaaaaaaa" - - -def test_a_late_write_lands_only_while_access_is_open(store: MemoryStore) -> None: - publish(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) - driver = MemoryDriver(store, controlled=True) - reader = Endpoint(driver, ["A", "B"], fill=0x2D) - op = driver.get(reader.rows) - driver.complete(op, [0]) - driver.fail(op, "timeout", [1]) - results = driver.poll(op) - - driver.write_late(op, 1, b"bbbbbbbb") - assert reader.read_row(1) == b"bbbbbbbb" - assert driver.poll(op) == results - assert not driver.access_ended(op) - - with pytest.raises(RuntimeError, match="ready, not faulted"): - driver.write_late(op, 0, b"cccccccc") - assert reader.read_row(0) == b"aaaaaaaa" - - driver.end_access(op) - with pytest.raises(RuntimeError, match="ended"): - driver.write_late(op, 1, b"cccccccc") - with pytest.raises(RuntimeError, match="ended"): - driver.transfer(op, [1]) - assert reader.read_row(1) == b"bbbbbbbb" - assert reader.intact() - - -def test_a_cancelled_row_cannot_be_written_late(store: MemoryStore) -> None: - driver = MemoryDriver(store, controlled=True) - reader = Endpoint(driver, ["A"], fill=0x2D) - op = driver.get(reader.rows) - driver.cancel(op) - assert statuses(driver, op) == [RowStatus.CANCELLED] - with pytest.raises(RuntimeError, match="cancelled, not faulted"): - driver.write_late(op, 0, b"aaaaaaaa") - assert reader.read_row(0) == b"-" * SLOT - - -def test_timeout_is_a_fault_not_a_miss(store: MemoryStore) -> None: - driver = MemoryDriver(store, controlled=True) - reader = Endpoint(driver, ["C"]) - op = driver.get(reader.rows) - driver.fail(op, "timeout") - (result,) = driver.poll(op) - assert result.status is RowStatus.FAULT - assert result.fault == "timeout" - - -class FailingStore(MemoryStore): - """Raises ``StoreError`` on the n-th commit, as a store that drops a connection would.""" - - def __init__(self, fail_on_commit: int) -> None: - super().__init__() - self._fail_on = fail_on_commit - self._commits = 0 - - def commit(self, key: str, data: bytes) -> bool: - self._commits += 1 - if self._commits == self._fail_on: - raise StoreError("connection dropped") - return super().commit(key, data) - - -def test_a_store_error_mid_operation_becomes_a_row_fault_and_keeps_the_handle() -> None: - store = FailingStore(fail_on_commit=2) - writer = Endpoint(MemoryDriver(store), ["A", "B", "C"]) - for i, data in enumerate((b"aaaaaaaa", b"bbbbbbbb", b"cccccccc")): - writer.fill_row(i, data) - op = writer.driver.put(writer.rows) - results = writer.driver.poll(op) - assert [r.status for r in results] == [RowStatus.READY, RowStatus.FAULT, RowStatus.READY] - assert results[1].fault == "connection dropped" - assert writer.driver.access_ended(op) - assert store.committed_payload(writer.rows[0].key) == b"aaaaaaaa" - assert not store.contains(writer.rows[1].key) - assert store.committed_payload(writer.rows[2].key) == b"cccccccc" - assert writer.driver.drain(0.0) is True - - -class GatedStore(MemoryStore): - """Holds every commit until the test opens the gate, and says when one is waiting.""" - - def __init__(self) -> None: - super().__init__() - self.gate = threading.Event() - self.waiting = threading.Event() - - def commit(self, key: str, data: bytes) -> bool: - self.waiting.set() - self.gate.wait() - return super().commit(key, data) - - -def test_queries_do_not_wait_on_another_operation_in_the_store() -> None: - store = GatedStore() - driver = MemoryDriver(store, controlled=True) - done = Endpoint(driver, ["A"]) - done.fill_row(0, b"aaaaaaaa") - finished = driver.get(done.rows) - driver.fail(finished, "timeout") - - slow = Endpoint(driver, ["B"]) - slow.fill_row(0, b"bbbbbbbb") - blocked = driver.put(slow.rows) - worker = threading.Thread(target=driver.transfer, args=(blocked,), daemon=True) - worker.start() - answered = threading.Event() - seen: dict[str, object] = {} - failure: list[Exception] = [] - - def query() -> None: - # Runs off the main thread: if a query waits on the store, the main thread can still - # open the gate and the test fails instead of hanging. - try: - seen["finished"] = statuses(driver, finished) - seen["blocked"] = statuses(driver, blocked) - seen["access_ended"] = driver.access_ended(blocked) - try: - driver.end_access(blocked) - except RuntimeError as error: - seen["end_access"] = str(error) - driver.cancel(blocked) - seen["after_cancel"] = statuses(driver, blocked) - except Exception as error: # surfaced on the main thread below - failure.append(error) - finally: - answered.set() - - querier = threading.Thread(target=query, daemon=True) - try: - assert store.waiting.wait(5.0), "the worker never reached the store" - # The worker is inside the store with its row taken; these must answer at once. - querier.start() - assert answered.wait(5.0), "queries waited on an operation that is inside the store" - finally: - store.gate.set() - worker.join(5.0) - if querier.ident is not None: # started; joining an unstarted thread raises - querier.join(5.0) - assert not worker.is_alive() and not querier.is_alive() - if failure: - raise AssertionError("a query raised") from failure[0] - assert seen["finished"] == [RowStatus.FAULT] - assert seen["blocked"] == [RowStatus.PENDING] - assert seen["access_ended"] is False - assert seen.get("end_access") == "a transfer is still in progress" - assert seen["after_cancel"] == [RowStatus.PENDING] - driver.publish_results(blocked) - assert statuses(driver, blocked) == [RowStatus.READY] - assert store.committed_payload(slow.rows[0].key) == b"bbbbbbbb" - - -# Cancellation, probe, drain. - - -def test_cancel_concludes_unstarted_rows_and_spares_started_ones(store: MemoryStore) -> None: - publish(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) - driver = MemoryDriver(store, controlled=True) - reader = Endpoint(driver, ["A", "B"], fill=0x2D) - op = driver.get(reader.rows) - driver.transfer(op, [0]) - driver.cancel(op) - assert statuses(driver, op) == [RowStatus.PENDING, RowStatus.CANCELLED] - driver.publish_results(op, [0]) - assert statuses(driver, op) == [RowStatus.READY, RowStatus.CANCELLED] - assert reader.read_row(1) == b"-" * SLOT - assert not driver.access_ended(op) - - -def test_an_ineffective_cancel_changes_nothing(store: MemoryStore) -> None: - driver = MemoryDriver(store, controlled=True) - driver.set_cancel_effective(False) - reader = Endpoint(driver, ["A"]) - op = driver.get(reader.rows) - driver.cancel(op) - assert statuses(driver, op) == [RowStatus.PENDING] - - -def test_probe_answers_hit_miss_or_unknown(store: MemoryStore) -> None: - publish(store, {"A": b"aaaaaaaa"}) - driver = MemoryDriver(store) - keys = [encode_key("t", NAMES["A"]), encode_key("t", NAMES["C"])] - assert driver.probe(keys, 1.0) == (ProbeStatus.HIT, ProbeStatus.MISS) - driver.set_probe_unknown(True) - assert driver.probe(keys, 1.0) == (ProbeStatus.UNKNOWN, ProbeStatus.UNKNOWN) - - -def test_drain_waits_for_access_end_and_refuses_new_work(store: MemoryStore) -> None: - publish(store, {"A": b"aaaaaaaa"}) - driver = MemoryDriver(store, controlled=True) - reader = Endpoint(driver, ["A"]) - op = driver.get(reader.rows) - assert driver.drain(0.01) is False - assert statuses(driver, op) == [RowStatus.CANCELLED] - with pytest.raises(ValueError, match="draining"): - driver.get(reader.rows) - - release = threading.Timer(0.05, driver.end_access, args=(op,)) - release.start() - try: - assert driver.drain(5.0) is True - finally: - release.join() - assert driver.open_ops() == [] - - -# Registration. - - -def test_register_checks_geometry_and_overlap_and_unregister_is_idempotent( - store: MemoryStore, -) -> None: - driver = MemoryDriver(store) - buffer = HostBuffer(64) - with pytest.raises(ValueError, match="layout needs"): - driver.register(buffer.address, 64, RegionLayout(SLOT, 9)) - handle = driver.register(buffer.address, 64, RegionLayout(SLOT, 8)) - with pytest.raises(ValueError, match="overlaps"): - driver.register(buffer.address + 8, 8, RegionLayout(8, 1)) - driver.unregister(handle) - driver.unregister(handle) - assert driver.live_registrations() == [] - driver.register(buffer.address + 8, 8, RegionLayout(8, 1)) - - -def test_a_registration_the_store_may_hold_keeps_its_handle(store: MemoryStore) -> None: - driver = MemoryDriver(store) - buffer = HostBuffer(64) - driver.fail_next_register("store unreachable", escaped=False) - with pytest.raises(RegistrationError) as clean: - driver.register(buffer.address, 64, RegionLayout(SLOT, 8)) - assert clean.value.handle is None - assert driver.live_registrations() == [] - - driver.fail_next_register("timed out after the request", escaped=True) - with pytest.raises(RegistrationError) as ambiguous: - driver.register(buffer.address, 64, RegionLayout(SLOT, 8)) - handle = ambiguous.value.handle - assert handle is not None - assert driver.live_registrations() == [handle] - - driver.fail_next_unregister(handle, "busy") - with pytest.raises(StoreError, match="busy"): - driver.unregister(handle) - assert driver.live_registrations() == [handle] - driver.unregister(handle) - assert driver.live_registrations() == [] - - -def test_memory_driver_has_every_member_of_the_driver_protocol(store: MemoryStore) -> None: - driver: StoreDriver = MemoryDriver(store) - assert isinstance(driver, StoreDriver) - - -def test_automatic_mode_follows_the_same_protocol(store: MemoryStore) -> None: - writer = Endpoint(MemoryDriver(store), ["A"]) - writer.fill_row(0, b"aaaaaaaa") - op = writer.driver.put(writer.rows) - assert concluded(writer.driver.poll(op)) - assert writer.driver.access_ended(op) - assert writer.driver.drain(0.0) is True - assert writer.driver.open_ops() == [] From dbf115a661b0e1fb966b42bb73bc6587bbc8d332 Mon Sep 17 00:00:00 2001 From: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> Date: Fri, 16 Oct 2026 10:56:18 +0800 Subject: [PATCH 3/6] [None][feat] Let get and probe carry an opaque hint for drivers that declare it The backend passes a hint only to a driver whose capabilities declare hints and rejects it before dispatch otherwise, so hint-free drivers keep their calls as they were. New capabilities default to unsupported; fields whose default would grant something stay explicit. Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> --- .../_torch/disaggregation/blob/store.py | 35 +++++++-- .../disaggregation/blob/driver_conformance.py | 6 +- .../disaggregation/blob/memory_store.py | 28 +++++-- .../disaggregation/blob/test_blob_format.py | 1 + .../disaggregation/blob/test_memory_driver.py | 76 +++++++++++++++++-- 5 files changed, 124 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/blob/store.py b/tensorrt_llm/_torch/disaggregation/blob/store.py index d2233e1aed75..97615d3eca24 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/store.py +++ b/tensorrt_llm/_torch/disaggregation/blob/store.py @@ -14,7 +14,12 @@ import enum from collections.abc import Sequence from dataclasses import dataclass -from typing import Protocol, runtime_checkable +from typing import Any, Protocol, runtime_checkable + +Hint = Any +"""The caller's hint for where a get or probe should look, handed to the driver unchanged. The +backend passes a hint only to a driver that declares ``hints`` and rejects it before dispatch for +any other, which therefore keeps its hint-free calls as they were.""" class StoreError(Exception): @@ -194,7 +199,11 @@ class ProbeStatus(enum.Enum): @dataclass(frozen=True) class DriverCapabilities: - """What a driver implements. Fixed for the driver's life; every field is declared. + """What a driver implements. Fixed for the driver's life. + + A field whose default would grant something, such as an unlimited row count, is declared + explicitly. A capability added later defaults to not supported, so existing drivers need no + change to stay honest. Args: get: ``get`` is implemented. @@ -202,6 +211,8 @@ class DriverCapabilities: cancel: ``cancel`` does something; without it the backend only suppresses unstarted work. probe: ``probe`` answers; without it the backend fetches blindly. max_rows_per_op: Rows one ``put``/``get`` accepts, or ``None`` for no limit. + hints: ``get`` and ``probe`` read a caller's hint; without it the backend rejects a hint + before dispatch and the driver is called without one. """ get: bool @@ -209,6 +220,7 @@ class DriverCapabilities: cancel: bool probe: bool max_rows_per_op: int | None + hints: bool = False def __post_init__(self) -> None: if not (self.get or self.put): @@ -277,8 +289,13 @@ def put(self, rows: Sequence[RowIO]) -> OpHandle: """ ... - def get(self, rows: Sequence[RowIO]) -> OpHandle: - """Fetch ``rows``; may return before any byte moved. Raises as ``put``.""" + def get(self, rows: Sequence[RowIO], hint: Hint = None) -> OpHandle: + """Fetch ``rows``, from where ``hint`` says when given; may return before any byte moved. + + Raises as ``put``. ``hint`` is passed only to a driver that declares ``hints``; the + backend rejects it before dispatch otherwise, and a driver may treat an unexpected hint + as a programming error. + """ ... def poll(self, op: OpHandle) -> tuple[RowResult, ...]: @@ -309,13 +326,17 @@ def cancel(self, op: OpHandle) -> None: Called only when ``capabilities.cancel`` is set.""" ... - def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...]: - """Advisory presence of ``keys``, answered within ``budget_s`` seconds. + def probe( + self, keys: Sequence[str], budget_s: float, hint: Hint = None + ) -> tuple[ProbeStatus, ...]: + """Advisory presence of ``keys``, at the source ``hint`` names when given, answered within + ``budget_s`` seconds. 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. Raises on a store error rather than answering MISS. + when ``capabilities.probe`` is set, and with a hint only when ``capabilities.hints`` is. + Raises on a store error rather than answering MISS. """ ... diff --git a/tests/unittest/_torch/disaggregation/blob/driver_conformance.py b/tests/unittest/_torch/disaggregation/blob/driver_conformance.py index 75275264886a..4faafe92a11e 100644 --- a/tests/unittest/_torch/disaggregation/blob/driver_conformance.py +++ b/tests/unittest/_torch/disaggregation/blob/driver_conformance.py @@ -27,6 +27,7 @@ from tensorrt_llm._torch.disaggregation.blob.keys import encode_key from tensorrt_llm._torch.disaggregation.blob.store import ( DriverRejected, + Hint, OpHandle, ProbeStatus, RegionHandle, @@ -153,8 +154,9 @@ def put(self, rows: Sequence[RowIO]) -> OpHandle: self.ops.append(op) return op - def get(self, rows: Sequence[RowIO]) -> OpHandle: - op = self.driver.get(rows) + def get(self, rows: Sequence[RowIO], hint: Hint = None) -> OpHandle: + # Without a hint the call keeps its old shape, as the backend's does. + op = self.driver.get(rows) if hint is None else self.driver.get(rows, hint) self.ops.append(op) return op diff --git a/tests/unittest/_torch/disaggregation/blob/memory_store.py b/tests/unittest/_torch/disaggregation/blob/memory_store.py index 1fb0c94bb2c7..f304d699c1db 100644 --- a/tests/unittest/_torch/disaggregation/blob/memory_store.py +++ b/tests/unittest/_torch/disaggregation/blob/memory_store.py @@ -37,6 +37,7 @@ READY, DriverCapabilities, DriverRejected, + Hint, OpHandle, ProbeStatus, RegionHandle, @@ -182,6 +183,7 @@ class MemoryOp: kind: Kind rows: list[_Row] + hint: Hint = None access_ended: bool = False unprovable: bool = False @@ -215,6 +217,7 @@ def __init__( self._probe_unknown = False self._register_failure: tuple[str, bool] | None = None self._unregister_failures: dict[int, str] = {} + self.hints_seen: list[Hint] = [] # every hint a get or probe was given, in order # The StoreDriver protocol. @@ -260,8 +263,8 @@ def unregister(self, handle: RegionHandle) -> None: def put(self, rows: Sequence[RowIO]) -> MemoryOp: return self._submit("put", rows) - def get(self, rows: Sequence[RowIO]) -> MemoryOp: - return self._submit("get", rows) + def get(self, rows: Sequence[RowIO], hint: Hint = None) -> MemoryOp: + return self._submit("get", rows, hint) def poll(self, op: OpHandle) -> tuple[RowResult, ...]: with self._lock: @@ -280,7 +283,14 @@ def cancel(self, op: OpHandle) -> None: if self._cancel_effective: self._cancel_unstarted(self._op(op)) - def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...]: + def probe( + self, keys: Sequence[str], budget_s: float, hint: Hint = None + ) -> tuple[ProbeStatus, ...]: + with self._lock: + self._require_not_closed() + self._check_hint(hint) + if hint is not None: + self.hints_seen.append(hint) answers: dict[str, ProbeStatus] = {} errors: list[Exception] = [] finished = threading.Event() @@ -299,7 +309,6 @@ def ask() -> None: # The store may block; the budget bounds this call, not the store. Registration and # start happen under the lock, so close and drain see every worker. with self._lock: - self._require_not_closed() self._probes = [probe for probe in self._probes if not probe.finished.is_set()] if self._draining or self._probe_unknown or len(self._probes) >= MAX_PROBES_IN_FLIGHT: return tuple(ProbeStatus.UNKNOWN for _ in keys) @@ -469,6 +478,10 @@ def _check_unique_keys(self, rows: Sequence[RowIO]) -> None: if len(set(keys)) != len(keys): raise DriverRejected("a key appears more than once in the operation") + def _check_hint(self, hint: Hint) -> None: + if hint is not None and not self._capabilities.hints: + raise DriverRejected("this driver reads no hints") + def _check_registered(self, rows: Sequence[RowIO]) -> None: for io in rows: if not self._registered(io.payload): @@ -504,16 +517,19 @@ def _quiet(self) -> bool: # Internals. - def _submit(self, kind: Kind, rows: Sequence[RowIO]) -> MemoryOp: + def _submit(self, kind: Kind, rows: Sequence[RowIO], hint: Hint = None) -> MemoryOp: with self._lock: self._require_accepting() self._check_direction(kind) self._check_not_draining() + self._check_hint(hint) self._check_row_limit(rows) self._check_unique_keys(rows) self._check_registered(rows) - op = MemoryOp(kind, [_Row(io) for io in rows]) + op = MemoryOp(kind, [_Row(io) for io in rows], hint) self._ops.append(op) + if hint is not None: + self.hints_seen.append(hint) if not self.controlled: self._before_execute(op) self._execute(op) diff --git a/tests/unittest/_torch/disaggregation/blob/test_blob_format.py b/tests/unittest/_torch/disaggregation/blob/test_blob_format.py index deb724150bf3..0bdb179640be 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_blob_format.py +++ b/tests/unittest/_torch/disaggregation/blob/test_blob_format.py @@ -75,6 +75,7 @@ def test_capabilities_are_declared_in_full() -> None: with pytest.raises(ValueError): DriverCapabilities(get=True, put=True, cancel=False, probe=False, max_rows_per_op=0) caps = DriverCapabilities(get=True, put=False, cancel=False, probe=False, max_rows_per_op=4) + assert caps.hints is False # a capability added later defaults to not supported with pytest.raises(AttributeError): caps.put = True # type: ignore[misc] diff --git a/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py index c84f5f49a6dc..39194107a188 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py +++ b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py @@ -36,6 +36,7 @@ MISS, DriverCapabilities, DriverRejected, + Hint, OpHandle, ProbeStatus, RegionHandle, @@ -129,6 +130,17 @@ def _run(target: Callable[..., object], *args: object) -> threading.Thread: return worker +class _HintFreeSignatures(MemoryDriver): + """A driver written before hints existed: ``get`` and ``probe`` take no hint. With + ``hints`` undeclared the backend never passes one, so it stays a valid driver.""" + + def get(self, rows: Sequence[RowIO]) -> MemoryOp: # type: ignore[override] + return super().get(rows) + + def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...]: # type: ignore[override] + return super().probe(keys, budget_s) + + class TestMemoryDriverConformance(DriverConformance): driver_cls: type[MemoryDriver] = MemoryDriver capabilities = DEFAULT_CAPABILITIES @@ -142,6 +154,10 @@ def make_blocked_probe(self) -> tuple[StoreDriver, str, Callable[[], None]]: return self.driver_cls(store, capabilities=self.capabilities), KEY_A, store.gate.set +class TestHintFreeDriverConformance(TestMemoryDriverConformance): + driver_cls = _HintFreeSignatures + + class TestMemoryDriver: """Behaviour only the in-memory driver's controls or stores can stage.""" @@ -560,6 +576,37 @@ def test_probes_in_flight_are_capped(self) -> None: self._shutdown(driver) assert driver.probes_in_flight() == 0 + def test_memory_driver_rejects_a_hint_it_cannot_read_before_any_effect(self) -> None: + """The backend rejects such hints before dispatch; this driver also refuses them itself, + so a backend bug cannot make it read with a hint it silently dropped.""" + driver = self.make() + self._publish(driver, {"A": b"aaaaaaaa"}) + reader = Endpoint(driver, ["A"], fill=0x2D) + with pytest.raises(DriverRejected): + driver.get(reader.rows(), {"source": "worker-7"}) + assert reader.read_row(0) == b"-" * SLOT + with pytest.raises(DriverRejected): + driver.probe([KEY_A], TIMEOUT, {"source": "worker-7"}) + assert statuses(driver, driver.get(reader.rows())) == [RowStatus.READY] + assert driver.hints_seen == [] + + def test_a_hint_reaches_a_driver_that_declares_hints_unchanged(self) -> None: + hinted = DriverCapabilities( + get=True, put=True, cancel=True, probe=True, max_rows_per_op=None, hints=True + ) + driver = self.make(capabilities=hinted) + self._publish(driver, {"A": b"aaaaaaaa"}) + reader = Endpoint(driver, ["A"], fill=0x2D) + hint = {"source": "worker-7"} + op = driver.get(reader.rows(), hint) + assert statuses(driver, op) == [RowStatus.READY] + assert reader.read_row(0) == b"aaaaaaaa" + assert driver.probe([KEY_A], TIMEOUT, hint) == (ProbeStatus.HIT,) + assert driver.hints_seen == [hint, hint] + assert all(seen is hint for seen in driver.hints_seen) # the same object, not a copy + assert statuses(driver, driver.get(reader.rows())) == [RowStatus.READY] + assert len(driver.hints_seen) == 2 # no hint, nothing recorded + def test_probe_unknown_switch(self) -> None: driver = self.make() self._publish(driver, {"A": b"aaaaaaaa"}) @@ -770,10 +817,10 @@ def _check_unique_keys(self, rows: Sequence[RowIO]) -> None: class _TouchesMemoryBeforeChecking(MemoryDriver): - def get(self, rows: Sequence[RowIO]) -> MemoryOp: + def get(self, rows: Sequence[RowIO], hint: Hint = None) -> MemoryOp: for io in rows: ctypes.memset(io.payload.address, 0, io.payload.nbytes) - return super().get(rows) + return super().get(rows, hint) class _MissTouchesTheDestination(MemoryDriver): @@ -827,9 +874,11 @@ def capabilities(self) -> DriverCapabilities: class _ProbeTakesItsTime(MemoryDriver): - def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...]: + def probe( + self, keys: Sequence[str], budget_s: float, hint: Hint = None + ) -> tuple[ProbeStatus, ...]: time.sleep(budget_s + 0.8) - return super().probe(keys, budget_s) + return super().probe(keys, budget_s, hint) class _KeepsAcceptingWhileDraining(MemoryDriver): @@ -892,15 +941,19 @@ def _put_row(self, io: RowIO) -> RowResult: class _ProbeIgnoresTheBudget(MemoryDriver): - def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...]: + def probe( + self, keys: Sequence[str], budget_s: float, hint: Hint = None + ) -> tuple[ProbeStatus, ...]: return tuple(self.store.presence(key) for key in keys) class _ProbesWithoutLimit(MemoryDriver): - def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...]: + def probe( + self, keys: Sequence[str], budget_s: float, hint: Hint = None + ) -> tuple[ProbeStatus, ...]: with self._lock: self._probes = self._probes[: MAX_PROBES_IN_FLIGHT - 1] - return super().probe(keys, budget_s) + return super().probe(keys, budget_s, hint) class _ReportsAccessEndedEarly(MemoryDriver): @@ -908,6 +961,11 @@ def access_ended(self, op: OpHandle) -> bool: return True +class _AcceptsHintsItCannotRead(MemoryDriver): + def _check_hint(self, hint: Hint) -> None: + pass + + class _ExecutesWithTheControls(MemoryDriver): def _execute(self, op: MemoryOp) -> None: self.transfer(op) @@ -1023,6 +1081,10 @@ def _drain_can_return(self) -> bool: _CallsDeclinedAFault, "test_a_row_the_store_has_no_room_for_is_declined_not_faulted", ), + "an unreadable hint is refused": ( + _AcceptsHintsItCannotRead, + "test_memory_driver_rejects_a_hint_it_cannot_read_before_any_effect", + ), "a submission racing a drain keeps its handle": ( _ExecutesWithTheControls, "test_a_drain_racing_a_submission_leaves_a_cancelled_operation_with_its_handle", From f59721d575b918f7416f9e3d40ca1a410e299483 Mon Sep 17 00:00:00 2001 From: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:39:41 +0800 Subject: [PATCH 4/6] [None][feat] Ship the store driver conformance suite in the wheel and widen rows to spans Rows are a key plus disjoint spans; DriverRejected stands apart from store errors; NOT_READY, DECLINED and CANCELLED are not faults, and new capabilities (absent_is_miss, device_memory, multi_span) default to unsupported. The conformance suite, harness, in-memory driver and broken drivers move to tensorrt_llm._torch.disaggregation.testing with scenario and memory hooks, per-harness key prefixes, and cleanup that retains memory it cannot prove released. Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> --- .../_torch/disaggregation/blob/store.py | 127 ++- .../_torch/disaggregation/testing/__init__.py | 13 + .../disaggregation/testing/broken_drivers.py | 315 +++++++ .../testing/driver_conformance.py | 842 ++++++++++++++++++ .../disaggregation/testing/host_memory.py | 77 ++ .../disaggregation/testing}/memory_store.py | 145 ++- .../disaggregation/testing}/object_trailer.py | 0 .../disaggregation/blob/driver_conformance.py | 505 ----------- .../_torch/disaggregation/blob/host_memory.py | 50 -- .../disaggregation/blob/test_blob_format.py | 47 +- .../disaggregation/blob/test_memory_driver.py | 769 ++++++++-------- .../blob/test_testing_isolation.py | 83 ++ 12 files changed, 1908 insertions(+), 1065 deletions(-) create mode 100644 tensorrt_llm/_torch/disaggregation/testing/__init__.py create mode 100644 tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py create mode 100644 tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py create mode 100644 tensorrt_llm/_torch/disaggregation/testing/host_memory.py rename {tests/unittest/_torch/disaggregation/blob => tensorrt_llm/_torch/disaggregation/testing}/memory_store.py (84%) rename {tests/unittest/_torch/disaggregation/blob => tensorrt_llm/_torch/disaggregation/testing}/object_trailer.py (100%) delete mode 100644 tests/unittest/_torch/disaggregation/blob/driver_conformance.py delete mode 100644 tests/unittest/_torch/disaggregation/blob/host_memory.py create mode 100644 tests/unittest/_torch/disaggregation/blob/test_testing_isolation.py diff --git a/tensorrt_llm/_torch/disaggregation/blob/store.py b/tensorrt_llm/_torch/disaggregation/blob/store.py index 97615d3eca24..90385b37496f 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/store.py +++ b/tensorrt_llm/_torch/disaggregation/blob/store.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """The private boundary between the blob backend and a byte store. -A driver moves whole rows between registered host memory and one store. It answers two questions +A driver moves whole rows between registered memory and one store. It answers two questions separately: what each row's logical result is (:meth:`StoreDriver.poll`), and whether an operation can still touch the caller's memory (:meth:`StoreDriver.access_ended`). Neither answer implies the other. Nothing here belongs to the public backend contract in ``base/``; the backend translates @@ -17,24 +17,27 @@ from typing import Any, Protocol, runtime_checkable Hint = Any -"""The caller's hint for where a get or probe should look, handed to the driver unchanged. The -backend passes a hint only to a driver that declares ``hints`` and rejects it before dispatch for -any other, which therefore keeps its hint-free calls as they were.""" +"""The caller's hint for where a get or probe should look, handed to the driver unchanged. A hint +is an immutable snapshot: the driver may keep a reference to it for as long as it needs and never +modifies it. The backend passes a hint only to a driver that declares ``hints``; for any other +driver it rejects the submission that carries one, before dispatch, so such a driver keeps its +hint-free calls as they were.""" -class StoreError(Exception): - """A store-side failure the driver could not express as a row result.""" - - -class DriverRejected(StoreError): +class DriverRejected(Exception): """Raised by a driver on purpose, before it touched the store or the caller's memory. - The only exception that certifies nothing escaped: the caller may forget the call. Any other + The only exception that certifies nothing escaped: the caller may forget the call. It is not a + :class:`StoreError`, so code that handles store trouble never swallows a rejection. Any other raise from ``register``, ``put`` or ``get`` may have escaped, and the caller keeps the memory protected until the driver proves otherwise, at the latest through a ``drain`` returning True. """ +class StoreError(Exception): + """A store-side failure the driver could not express as a row result.""" + + class RegistrationError(StoreError): """Registering a region failed after the store was asked. @@ -52,9 +55,9 @@ def __init__(self, reason: str, handle: RegionHandle | None = None) -> None: 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 and DECLINED say that - nothing moved and nothing went wrong, so the backend reports the row as not delivered without - a fault and they never count toward isolation. + 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. """ PENDING = "pending" @@ -66,7 +69,8 @@ class RowStatus(enum.Enum): store's "already exists" counts only once that object is committed.""" MISS = "miss" - """For a get: the store holds no object under the key.""" + """For a get: the store holds no object under the key. Only a driver that declares + ``absent_is_miss`` can tell this apart from a failure; another reports FAULT for such a key.""" NOT_READY = "not_ready" """The key is held by a write that has not finished: a get delivered nothing, a put published @@ -77,10 +81,12 @@ class RowStatus(enum.Enum): the other rows of the operation go on.""" CANCELLED = "cancelled" - """Never started, because of a cancel or a drain; the row's memory was not touched.""" + """Never started, because of a cancel or a drain; no byte moved and the row's memory was not + touched. A row that had started never concludes CANCELLED.""" FAULT = "fault" - """Anything else: timeout, size or integrity mismatch, transport or store error.""" + """Anything else: timeout, size or integrity mismatch, transport or store error. Says nothing + about the destination's contents.""" @dataclass(frozen=True) @@ -99,14 +105,6 @@ def concluded(self) -> bool: return self.status is not RowStatus.PENDING -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) - - def faulted(reason: str) -> RowResult: """A FAULT result carrying ``reason``.""" return RowResult(RowStatus.FAULT, reason) @@ -119,7 +117,11 @@ def concluded(results: Sequence[RowResult]) -> bool: @dataclass(frozen=True) class Span: - """Host memory ``[address, address + nbytes)``.""" + """A non-empty range ``[address, address + nbytes)`` inside a registered region. + + Whether that region is host or device memory is the region's business, not the span's. A span + need not be a whole slot. + """ address: int nbytes: int @@ -163,21 +165,34 @@ def nbytes(self) -> int: @dataclass(frozen=True) class RowIO: - """One row to move: its store key and where its payload sits. + """One row to move: its store key and the spans its payload occupies. - The payload is the KV bytes of one staging slot in a registered region; the object a driver - stores is that payload, in whatever format and with whatever integrity data the driver keeps. - ``position`` is ``(group, row)`` within the submission, kept so the backend can map results - back; drivers treat it as opaque. + The stored object is the spans' bytes back to back, in the order given; the key does not + depend on how a row is split. A row staged on the host has one span. A row lent in place has + one span per pool of its pool group, in pool order, and may cover only part of a slot when the + two sides of a transfer are sharded differently. Spans of one row never overlap. A row of more + than one span reaches only a driver that declares ``multi_span``. """ key: str - payload: Span - position: tuple[int, int] | None = None + spans: tuple[Span, ...] def __post_init__(self) -> None: if not self.key: raise ValueError("a row needs a store key") + spans = tuple(self.spans) + object.__setattr__(self, "spans", spans) + if not spans: + raise ValueError(f"row {self.key} has no spans") + for i, first in enumerate(spans): + for second in spans[i + 1 :]: + if first.overlaps(second): + raise ValueError(f"row {self.key} has overlapping spans {first} and {second}") + + @property + def nbytes(self) -> int: + """The object's size: the spans' bytes added up.""" + return sum(span.nbytes for span in self.spans) class RegionHandle(Protocol): @@ -211,8 +226,15 @@ class DriverCapabilities: cancel: ``cancel`` does something; without it the backend only suppresses unstarted work. probe: ``probe`` answers; without it the backend fetches blindly. max_rows_per_op: Rows one ``put``/``get`` accepts, or ``None`` for no limit. - hints: ``get`` and ``probe`` read a caller's hint; without it the backend rejects a hint - before dispatch and the driver is called without one. + hints: ``get`` and ``probe`` read a caller's hint; without it the backend rejects a + submission carrying a hint before dispatch and the driver is called without one. + absent_is_miss: A get of a key the store does not hold concludes MISS. Without it the + store cannot tell absence from failure, such a get may conclude FAULT, and the + backend's fault accounting counts misses too. + device_memory: The driver can register regions of device memory; without it every region + handed to it is host memory. + multi_span: ``put`` and ``get`` take rows of more than one span; without it every row + handed to it has exactly one. """ get: bool @@ -221,6 +243,9 @@ class DriverCapabilities: probe: bool max_rows_per_op: int | None hints: bool = False + absent_is_miss: bool = False + device_memory: bool = False + multi_span: bool = False def __post_init__(self) -> None: if not (self.get or self.put): @@ -245,14 +270,19 @@ class StoreDriver(Protocol): rows' memory until ``access_ended`` is true. Neither a failure result nor a successful ``unregister`` is access-end evidence. + Memory: the caller keeps the spans of a put unchanged, and lets nothing but the operation itself + write the spans of a get, until ``access_ended`` is true; a logical result does not lift that. The + driver never reads or writes outside the spans of the rows it was given. + Integrity: a READY get row is the whole object, intact, under the right key; a READY put row is a complete object any reader can get. How a driver assures this is its own: a store's native size checks, hidden unfinished objects and checksums, or data the driver keeps beside the object. A failed, declined or cancelled put never damages an existing complete object. - Queries (``poll``, ``access_ended``, ``access_unprovable``, ``cancel``) read recorded state: - they never block, never ask the store and never raise for a handle this driver issued. A handle - from elsewhere, or any call after ``close``, is a programming error and raises. + Non-blocking calls (``poll``, ``access_ended``, ``access_unprovable``, ``cancel``) read or + update recorded state: they never block, never ask the store and never raise for a handle this + driver issued. A handle from elsewhere, or any call after ``close``, is a programming error + and raises. """ @property @@ -261,7 +291,10 @@ def capabilities(self) -> DriverCapabilities: ... def register(self, address: int, nbytes: int, layout: RegionLayout) -> RegionHandle: - """Register one host region the driver may read and write. + """Register one region the driver may read and write. + + The region is host memory unless the driver declares ``device_memory``, in which case the + caller may also hand it device regions. Raises: DriverRejected: Bad geometry or overlap with a live registration; nothing escaped. @@ -282,19 +315,27 @@ def unregister(self, handle: RegionHandle) -> None: def put(self, rows: Sequence[RowIO]) -> OpHandle: """Publish ``rows``; may return before any byte moved. + The caller guarantees that the rows' spans hold the payload and stay unchanged until + ``access_ended`` is true. + Raises: DriverRejected: ``put`` is not declared, the driver is draining or closed, a key - appears twice in ``rows``, ``rows`` exceed ``max_rows_per_op``, or a payload lies - outside a live registration; nothing escaped. + appears twice in ``rows``, ``rows`` exceed ``max_rows_per_op``, a row has more + than one span without ``multi_span``, or a span lies outside a live registration; + nothing escaped. """ ... def get(self, rows: Sequence[RowIO], hint: Hint = None) -> OpHandle: """Fetch ``rows``, from where ``hint`` says when given; may return before any byte moved. - Raises as ``put``. ``hint`` is passed only to a driver that declares ``hints``; the - backend rejects it before dispatch otherwise, and a driver may treat an unexpected hint - as a programming error. + The caller guarantees that no two spans of one call overlap, within a row or across rows, + so a driver may write the rows concurrently without checking, and that nothing but this + operation writes those spans until ``access_ended`` is true. ``hint`` is passed only to a + driver that declares ``hints``; the backend rejects the submission otherwise, and a driver + may treat an unexpected hint as a programming error. + + Raises as ``put``. """ ... diff --git a/tensorrt_llm/_torch/disaggregation/testing/__init__.py b/tensorrt_llm/_torch/disaggregation/testing/__init__.py new file mode 100644 index 000000000000..511decccb36a --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/testing/__init__.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Test tooling for store drivers, shipped so a driver's own CI can import it from the wheel. + +``driver_conformance`` holds the checks every ``StoreDriver`` must pass and the harness that runs +them; ``memory_store`` holds the in-process store and driver that stand in for a real store; +``broken_drivers`` holds drivers that each break one rule, proving the checks bite; +``host_memory`` and ``object_trailer`` are their helpers. + +This package depends on pytest and is imported by tests only. No engine module imports it, and a +test guards that. It provides no public configuration: the in-memory driver cannot be selected +in a deployment. +""" diff --git a/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py b/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py new file mode 100644 index 000000000000..cbe1a25550bc --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py @@ -0,0 +1,315 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Drivers that each break exactly one rule of the driver boundary. + +Each is paired with the conformance check that guards its rule. A test runs the check against the +sound in-memory driver, which must pass, and against the broken one, which must fail; a check that +a weakened assertion lets both pass is caught there, before a real driver meets it. +""" + +from __future__ import annotations + +import ctypes +import time +from collections.abc import Sequence + +from tensorrt_llm._torch.disaggregation.blob.store import ( + DriverCapabilities, + DriverRejected, + Hint, + ProbeStatus, + RegionHandle, + RowIO, + RowResult, + RowStatus, + StoreError, + faulted, +) + +from .memory_store import DEFAULT_CAPABILITIES, Kind, MemoryDriver, MemoryOp, _scatter + +_MISS = RowResult(RowStatus.MISS) +_READY = RowResult(RowStatus.READY) + + +class AcceptsDuplicateKeys(MemoryDriver): + def _check_unique_keys(self, rows: Sequence[RowIO]) -> None: + pass + + +class TouchesMemoryBeforeChecking(MemoryDriver): + def get(self, rows: Sequence[RowIO], hint: Hint = None) -> MemoryOp: + for io in rows: + for span in io.spans: + ctypes.memset(span.address, 0, span.nbytes) + return super().get(rows, hint) + + +class MissTouchesTheDestination(MemoryDriver): + def _get_row(self, io: RowIO) -> RowResult: + result = super()._get_row(io) + if result.status is RowStatus.MISS: + _scatter(io.spans, bytes(io.nbytes)) + return result + + +class TruncatesLargerObjects(MemoryDriver): + def _get_row(self, io: RowIO) -> RowResult: + stored = self.store.get(io.key) + if stored is not None and len(stored.payload) > io.nbytes: + _scatter(io.spans, stored.payload[: io.nbytes]) + return _READY + return super()._get_row(io) + + +class CallsAWrongSizeAMiss(MemoryDriver): + def _get_row(self, io: RowIO) -> RowResult: + result = super()._get_row(io) + return _MISS if result.status is RowStatus.FAULT else result + + +class OverwritesExistingObjects(MemoryDriver): + def _put_row(self, io: RowIO) -> RowResult: + payload = b"".join(ctypes.string_at(span.address, span.nbytes) for span in io.spans) + self.store.overwrite(io.key, payload, self._trailer_for(io, payload)) + return super()._put_row(io) + + +class WritesTheWholeRegion(MemoryDriver): + def _get_row(self, io: RowIO) -> RowResult: + result = super()._get_row(io) + for registration in self.live_registrations(): + if io.spans[0].within(registration.span): + ctypes.memset(registration.span.address, 0, registration.span.nbytes) + return result + + +class OrdersSpansByAddress(MemoryDriver): + def _put_row(self, io: RowIO) -> RowResult: + return super()._put_row(RowIO(io.key, tuple(sorted(io.spans, key=lambda s: s.address)))) + + def _get_row(self, io: RowIO) -> RowResult: + return super()._get_row(RowIO(io.key, tuple(sorted(io.spans, key=lambda s: s.address)))) + + +class IgnoresTheRowLimit(MemoryDriver): + def _check_row_limit(self, rows: Sequence[RowIO]) -> None: + if not rows: + raise DriverRejected("an operation needs at least one row") + + +class IgnoresTheDirection(MemoryDriver): + def _check_direction(self, kind: Kind) -> None: + pass + + +class NeverEndsAccess(MemoryDriver): + def _execute(self, op: MemoryOp) -> None: + super()._execute(op) + with self._lock: + op.access_ended = False # the evidence is lost again + + +class ChangesItsCapabilities(MemoryDriver): + _flip = False + + @property + def capabilities(self) -> DriverCapabilities: + self._flip = not self._flip + return DriverCapabilities( + get=True, + put=True, + cancel=self._flip, + probe=True, + max_rows_per_op=None, + absent_is_miss=True, + multi_span=True, + ) + + +class ProbeTakesItsTime(MemoryDriver): + def probe( + self, keys: Sequence[str], budget_s: float, hint: Hint = None + ) -> tuple[ProbeStatus, ...]: + time.sleep(budget_s + 0.8) + return super().probe(keys, budget_s, hint) + + +class ProbeIgnoresTheBudget(MemoryDriver): + def probe( + self, keys: Sequence[str], budget_s: float, hint: Hint = None + ) -> tuple[ProbeStatus, ...]: + return tuple(self.store.presence(key) for key in keys) + + +class ProbeHidesStoreErrors(MemoryDriver): + def probe( + self, keys: Sequence[str], budget_s: float, hint: Hint = None + ) -> tuple[ProbeStatus, ...]: + try: + return super().probe(keys, budget_s, hint) + except StoreError: + return tuple(ProbeStatus.MISS for _ in keys) + + +class HoldsTheLockAcrossTheStore(MemoryDriver): + def _execute(self, op: MemoryOp) -> None: + with self._lock: + super()._execute(op) + + +class CallsNotReadyAFault(MemoryDriver): + def _get_row(self, io: RowIO) -> RowResult: + result = super()._get_row(io) + return faulted("not ready") if result.status is RowStatus.NOT_READY else result + + +class CallsDeclinedAFault(MemoryDriver): + def _put_row(self, io: RowIO) -> RowResult: + result = super()._put_row(io) + return faulted("no space") if result.status is RowStatus.DECLINED else result + + +class LetsRaisesEscape(MemoryDriver): + def _move(self, kind: Kind, io: RowIO) -> RowResult: + try: + return self._put_row(io) if kind == "put" else self._get_row(io) + except StoreError as error: + return faulted(str(error)) + + +class KeepsAcceptingWhileDraining(MemoryDriver): + def _check_not_draining(self) -> None: + pass + + +class ClosesAnyTime(MemoryDriver): + def _check_closable(self) -> None: + pass + + +class UnregisterKeepsTheRange(MemoryDriver): + def unregister(self, handle: RegionHandle) -> None: + pass + + +LIMITED = DriverCapabilities( + get=True, put=True, cancel=True, probe=True, max_rows_per_op=1, absent_is_miss=True +) +GET_ONLY = DriverCapabilities( + get=True, put=False, cancel=True, probe=True, max_rows_per_op=None, absent_is_miss=True +) + +# rule -> (driver breaking it, the capabilities to run with, the conformance check that must fail) +BROKEN_CONFORMANCE: dict[str, tuple[type[MemoryDriver], DriverCapabilities, str]] = { + "a key appears once per operation": ( + AcceptsDuplicateKeys, + DEFAULT_CAPABILITIES, + "test_a_key_appears_at_most_once_per_operation", + ), + "rejection happens before any effect": ( + TouchesMemoryBeforeChecking, + DEFAULT_CAPABILITIES, + "test_rows_outside_a_live_registration_are_rejected_before_any_effect", + ), + "a miss leaves the destination alone": ( + MissTouchesTheDestination, + DEFAULT_CAPABILITIES, + "test_get_reports_hit_hit_miss_and_leaves_the_miss_untouched", + ), + "a larger object is not truncated": ( + TruncatesLargerObjects, + DEFAULT_CAPABILITIES, + "test_an_object_larger_than_its_row_is_a_fault_not_a_truncated_ready", + ), + "a smaller object is a fault": ( + CallsAWrongSizeAMiss, + DEFAULT_CAPABILITIES, + "test_an_object_smaller_than_its_row_is_a_fault", + ), + "a put never damages an existing object": ( + OverwritesExistingObjects, + DEFAULT_CAPABILITIES, + "test_a_put_under_an_existing_key_leaves_the_object_whole", + ), + "an operation stays within its rows": ( + WritesTheWholeRegion, + DEFAULT_CAPABILITIES, + "test_an_operation_leaves_the_rest_of_its_region_alone", + ), + "spans keep the order given": ( + OrdersSpansByAddress, + DEFAULT_CAPABILITIES, + "test_a_multi_span_row_is_stored_and_fetched_back_to_back", + ), + "the row limit is enforced": ( + IgnoresTheRowLimit, + LIMITED, + "test_rows_beyond_the_declared_limit_are_rejected", + ), + "an undeclared direction is rejected": ( + IgnoresTheDirection, + GET_ONLY, + "test_an_undeclared_direction_is_rejected", + ), + "access ends with evidence": ( + NeverEndsAccess, + DEFAULT_CAPABILITIES, + "test_results_are_stable_and_access_ends_with_evidence", + ), + "capabilities are fixed": ( + ChangesItsCapabilities, + DEFAULT_CAPABILITIES, + "test_capabilities_do_not_change", + ), + "probe returns within its budget": ( + ProbeTakesItsTime, + DEFAULT_CAPABILITIES, + "test_probe_returns_within_its_budget", + ), + "probe answers unknown for a blocked key": ( + ProbeIgnoresTheBudget, + DEFAULT_CAPABILITIES, + "test_probe_answers_unknown_for_a_key_the_store_does_not_answer_in_time", + ), + "probe raises on a store error": ( + ProbeHidesStoreErrors, + DEFAULT_CAPABILITIES, + "test_probe_raises_on_a_store_error_rather_than_answering", + ), + "a blocked get delays nothing else": ( + HoldsTheLockAcrossTheStore, + DEFAULT_CAPABILITIES, + "test_a_blocked_get_delays_neither_another_get_nor_the_non_blocking_calls", + ), + "an unfinished write is not ready": ( + CallsNotReadyAFault, + DEFAULT_CAPABILITIES, + "test_a_key_held_by_an_unfinished_write_is_not_ready_for_get_put_and_probe", + ), + "no room is declined": ( + CallsDeclinedAFault, + DEFAULT_CAPABILITIES, + "test_a_row_with_no_room_is_declined_while_the_others_go_on", + ), + "a recorded operation never raises": ( + LetsRaisesEscape, + DEFAULT_CAPABILITIES, + "test_a_raise_after_the_operation_is_recorded_becomes_a_row_fault", + ), + "drain is one-way": ( + KeepsAcceptingWhileDraining, + DEFAULT_CAPABILITIES, + "test_drain_is_one_way", + ), + "close needs a true drain": ( + ClosesAnyTime, + DEFAULT_CAPABILITIES, + "test_close_needs_a_true_drain_and_no_live_registrations", + ), + "unregister frees the range": ( + UnregisterKeepsTheRange, + DEFAULT_CAPABILITIES, + "test_unregister_is_idempotent_and_frees_the_range", + ), +} diff --git a/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py b/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py new file mode 100644 index 000000000000..9d4aacb92f92 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py @@ -0,0 +1,842 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""What any ``StoreDriver`` owes the backend, checked through the protocol alone. + +A driver's test module subclasses :class:`DriverConformance` as ``Test...``, provides +``make_driver`` and, where its store can stage them, the scenario hooks (a blocked query, a key +another writer holds, a store with no room, and so on); a hook that returns ``None`` skips its +check. Every check observes consequences the backend could observe: bytes in registered memory, +row results, access end, rejections without effect. Nothing here uses a driver's own controls, +so the same checks run against the in-memory driver, a real store's driver, and the deliberately +broken drivers that prove each check bites. + +Each check runs inside a :class:`Harness` that keeps every endpoint and operation it created +under a key prefix of its own and, when the check is over, waits for access to end, unregisters +what it registered, drains and shuts the driver down. Memory whose access never provably ended, +or that the driver still holds, is kept alive instead of being freed. +""" + +from __future__ import annotations + +import contextlib +import hashlib +import secrets +import threading +import time +from collections.abc import Callable, Generator, Sequence + +import pytest + +from tensorrt_llm._torch.disaggregation.blob.keys import encode_key +from tensorrt_llm._torch.disaggregation.blob.store import ( + DriverRejected, + Hint, + OpHandle, + ProbeStatus, + RegionHandle, + RegionLayout, + RegistrationError, + RowIO, + RowResult, + RowStatus, + Span, + StoreDriver, + concluded, +) + +from .host_memory import HostBuffer, Memory + +SLOT = 8 + +# Lender names are 54 bytes and start with a 16-byte scope digest shared by every name of a +# deployment. A key built from a truncated name collides across these. +SCOPE = b"scope" + bytes(11) +NAME_BYTES = 54 + + +def name_of(label: str) -> bytes: + return SCOPE + hashlib.blake2b(label.encode(), digest_size=NAME_BYTES - len(SCOPE)).digest() + + +NAMES = {label: name_of(label) for label in "ABC"} + + +def key_of(label: str, prefix: str = "t") -> str: + return encode_key(prefix, name_of(label)) + + +def wait_until(predicate: Callable[[], bool], what: str, timeout: float) -> None: + """Polls ``predicate`` until true; fails after ``timeout`` seconds.""" + deadline = time.monotonic() + timeout + while not predicate(): + assert time.monotonic() < deadline, f"timed out waiting for {what}" + time.sleep(0.001) + + +def statuses(driver: StoreDriver, op: OpHandle) -> list[RowStatus]: + return [result.status for result in driver.poll(op)] + + +Allocation = tuple[Memory, "RegionHandle | None"] +Unblock = Callable[[], None] + + +class Endpoint: + """One side of a transfer: one region with a slot per label, under a key prefix.""" + + def __init__( + self, + driver: StoreDriver, + labels: Sequence[str], + memory: Memory, + handle: RegionHandle | None, + *, + prefix: str, + slot: int = SLOT, + ) -> None: + self.driver = driver + self.labels = list(labels) + self.memory = memory + self.prefix = prefix + self.slot = slot + self.handle = handle + # A handle given up front belongs to whoever allocated the memory; this endpoint only + # unregisters what it registered itself. + self.owns_registration = handle is None + self.released = False # unregistered or given back; a cleanup retry skips it + + def register(self) -> None: + """Registers the region. When the store may have kept it despite an error, the handle + is recorded all the same, so the region can still be unregistered later.""" + layout = RegionLayout(self.slot, len(self.labels), "kv") + try: + self.handle = self.driver.register(self.memory.address, self.memory.nbytes, layout) + except RegistrationError as error: + self.handle = error.handle + raise + + def key(self, label: str) -> str: + return key_of(label, self.prefix) + + def rows(self) -> list[RowIO]: + return [RowIO(self.key(label), (self.span(i),)) for i, label in enumerate(self.labels)] + + def span_at(self, offset: int, nbytes: int) -> Span: + """Any range of the region, slot boundaries or not.""" + return Span(self.memory.address + offset, nbytes) + + def span(self, index: int, nbytes: int | None = None) -> Span: + return Span( + self.memory.address + index * self.slot, self.slot if nbytes is None else nbytes + ) + + def fill_row(self, index: int, data: bytes) -> None: + self.memory.write(data, index * self.slot) + + def read_row(self, index: int) -> bytes: + return self.memory.read(index * self.slot, self.slot) + + def intact(self) -> bool: + return self.memory.canaries_intact() + + def unregister(self) -> None: + if self.owns_registration and self.handle is not None: + self.driver.unregister(self.handle) + + +# Endpoints whose memory a driver may still touch or still holds registered when a check ends. +# Never freed. +UNRELEASED_ENDPOINTS: list[Endpoint] = [] + + +class Harness: + """One driver under test plus everything a check created with it.""" + + def __init__(self, suite: DriverConformance, driver: StoreDriver, timeout: float) -> None: + self.suite = suite + self.driver = driver + self.timeout = timeout + self.prefix = "t" + secrets.token_hex(4) # keys of one check never meet another's + self.endpoints: list[Endpoint] = [] + self.ops: list[OpHandle] = [] + self.done = False + + # Creating things the harness will clean up. + + def key(self, label: str) -> str: + return key_of(label, self.prefix) + + def endpoint(self, labels: Sequence[str], *, fill: int = 0, slot: int = SLOT) -> Endpoint: + memory, handle = self.suite.allocate(slot * len(labels), fill) + endpoint = Endpoint(self.driver, labels, memory, handle, prefix=self.prefix, slot=slot) + self.endpoints.append(endpoint) # kept before the store is asked, whatever register does + if handle is None: + endpoint.register() + return endpoint + + def put(self, rows: Sequence[RowIO]) -> OpHandle: + op = self.driver.put(rows) + self.ops.append(op) + return op + + def get(self, rows: Sequence[RowIO], hint: Hint = None) -> OpHandle: + # Without a hint the call keeps its old shape, as the backend's does. + op = self.driver.get(rows) if hint is None else self.driver.get(rows, hint) + self.ops.append(op) + return op + + def publish(self, rows: dict[str, bytes], slot: int = SLOT) -> Endpoint: + """Puts ``rows`` (label to payload) and returns the writer endpoint.""" + writer = self.endpoint(list(rows), slot=slot) + for i, data in enumerate(rows.values()): + writer.fill_row(i, data) + op = self.put(writer.rows()) + assert [r.status for r in self.settle(op)] == [RowStatus.READY] * len(rows) + self.settle_access(op) + return writer + + # Waiting and asserting. + + def settle(self, op: OpHandle) -> tuple[RowResult, ...]: + """Waits for every row's result, not for access end.""" + wait_until( + lambda: concluded(self.driver.poll(op)), "the operation to conclude", self.timeout + ) + return self.driver.poll(op) + + def settle_access(self, op: OpHandle) -> None: + wait_until(lambda: self.driver.access_ended(op), "access to end", self.timeout) + + def assert_absent(self, result: RowResult, endpoint: Endpoint, index: int) -> None: + """What a get of a key nobody published must look like: MISS with the slot untouched + for a driver that can tell absence apart, otherwise anything but READY.""" + if self.driver.capabilities.absent_is_miss: + assert result.status is RowStatus.MISS, result + assert endpoint.read_row(index) == bytes([0x2D]) * endpoint.slot + else: + assert result.status is not RowStatus.READY, result + + def needs( + self, + *, + put: bool = False, + get: bool = False, + probe: bool = False, + cancel: bool = False, + multi_span: bool = False, + registers: bool = False, + rows: int = 1, + ) -> None: + """Skips the check when the driver does not declare, or the suite cannot stage, what + the check needs.""" + caps = self.driver.capabilities + if put and not caps.put: + pytest.skip("this driver declares no put") + if get and not caps.get: + pytest.skip("this driver declares no get") + if probe and not caps.probe: + pytest.skip("this driver declares no probe") + if cancel and not caps.cancel: + pytest.skip("this driver declares no cancel") + if multi_span and not caps.multi_span: + pytest.skip("this driver takes rows of one span") + if registers and not self.suite.registers_on_demand: + pytest.skip("this driver's memory is registered up front") + if caps.max_rows_per_op is not None and rows > caps.max_rows_per_op: + pytest.skip(f"this driver takes at most {caps.max_rows_per_op} rows per operation") + + # Cleanup. + + def cleanup(self) -> list[str]: + """Ends everything the check created, in order; returns what did not go to plan. + + Whenever cleanup stops short of shutting the driver down, every endpoint is retained for + good: memory the driver may still touch is never freed. An operation whose end the driver + declares unprovable is allowed by the contract and reported as no problem; one that + neither ends nor says so is. A cleanup that did not finish may be called again. + """ + if self.done: + return [] + problems: list[str] = [] + unprovable = False + for op in self.ops: + try: + wait_until( + lambda: self.driver.access_ended(op) or self._unprovable(op), + "access to end", + self.timeout, + ) + if not self.driver.access_ended(op): + unprovable = True + except AssertionError as error: + problems.append(str(error)) + except Exception as error: # noqa: BLE001 - every cleanup step is reported + problems.append(f"access end: {error!r}") + if problems or unprovable: + self._retain() + return problems + for endpoint in self.endpoints: + if endpoint.released: + continue + try: + if endpoint.owns_registration: + endpoint.unregister() + else: + self.suite.release(endpoint.memory) + endpoint.released = True + except Exception as error: # noqa: BLE001 - every cleanup step is reported + problems.append(f"unregister: {error!r}") + if problems: + self._retain() + return problems + try: + if self.driver.drain(self.timeout): + self.suite.shutdown(self.driver) + else: + problems.append("drain did not return True") + except Exception as error: # noqa: BLE001 - every cleanup step is reported + problems.append(f"drain/shutdown: {error!r}") + if problems: + self._retain() + return problems + self.done = True + return [] + + def _unprovable(self, op: OpHandle) -> bool: + try: + return bool(self.driver.access_unprovable(op)) + except Exception: # noqa: BLE001 - a raising query is not evidence either way + return False + + def _retain(self) -> None: + for endpoint in self.endpoints: + if endpoint not in UNRELEASED_ENDPOINTS: + UNRELEASED_ENDPOINTS.append(endpoint) + + +class DriverConformance: + """Contract checks for one driver. Subclass as ``Test...``, implement ``make_driver`` and + whichever scenario hooks the driver's store can stage.""" + + timeout = 5.0 + registers_on_demand = True # endpoints register fresh memory; a pool driver sets False + + def make_driver(self) -> StoreDriver: + raise NotImplementedError + + # Memory and lifetime hooks. + + def allocate(self, nbytes: int, fill: int) -> Allocation: + """Memory for one endpoint and, when the driver registered it up front, its handle. + + The memory object keeps its backing allocation alive for as long as it exists: a check + that cannot prove access ended retains the endpoint, and nothing else, for good. + """ + return HostBuffer(nbytes, fill), None + + def release(self, memory: Memory) -> None: + """Gives back memory that ``allocate`` handed out with a handle; the suite never + unregisters such memory itself.""" + + def shutdown(self, driver: StoreDriver) -> None: + """Ends the driver once it is drained; a pool driver unregisters its pool first.""" + driver.close() + + # Scenario hooks; ``None`` skips the check that needs it. + + def make_blocked_probe(self) -> tuple[StoreDriver, str, Unblock] | None: + """A driver whose store stalls every query of the returned key until the callable runs.""" + return None + + def make_blocked_get(self) -> tuple[StoreDriver, str, threading.Event | None, Unblock] | None: + """A driver whose store stalls every get of the returned key until the callable runs, + with an event set once a get is stalled, or ``None`` for no such signal.""" + return None + + def make_held_key(self) -> tuple[StoreDriver, str, Unblock] | None: + """A driver whose store holds the returned key as another writer's unfinished object + until the callable runs.""" + return None + + def make_full_store(self) -> StoreDriver | None: + """A driver whose store has room for exactly one more object.""" + return None + + def make_crash_after_record(self) -> StoreDriver | None: + """A driver whose store raises something that is not a store error from inside a put.""" + return None + + def make_failing_presence(self) -> tuple[StoreDriver, str] | None: + """A driver whose store fails the next presence query of the returned key.""" + return None + + @contextlib.contextmanager + def harness(self, driver: StoreDriver | None = None) -> Generator[Harness, None, None]: + """A fresh driver (or ``driver``) with cleanup afterwards. A failing check keeps its own + error; cleanup trouble fails only a check that passed.""" + h = Harness(self, driver if driver is not None else self.make_driver(), self.timeout) + try: + yield h + except BaseException: + h.cleanup() + raise + problems = h.cleanup() + assert not problems, "cleanup after the check: " + "; ".join(problems) + + # Bytes and results. + + def test_put_then_get_moves_whole_rows_and_never_aliases_the_source(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True, rows=2) + writer = h.publish({"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + writer.fill_row(0, b"xxxxxxxx") # the source changes after the put concluded + + reader = h.endpoint(["A", "B"], fill=0x2D) + results = h.settle(h.get(reader.rows())) + assert [r.status for r in results] == [RowStatus.READY, RowStatus.READY] + assert reader.read_row(0) == b"aaaaaaaa" + assert reader.read_row(1) == b"bbbbbbbb" + assert reader.intact() and writer.intact() + + def test_get_reports_hit_hit_miss_and_leaves_the_miss_untouched(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True, rows=3) + h.publish({"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + reader = h.endpoint(["A", "B", "C"], fill=0x2D) + results = h.settle(h.get(reader.rows())) + assert [r.status for r in results[:2]] == [RowStatus.READY, RowStatus.READY] + h.assert_absent(results[2], reader, 2) + assert reader.intact() + + def test_an_object_larger_than_its_row_is_a_fault_not_a_truncated_ready(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True) + h.publish({"A": b"aaaaaaaa"}, slot=8) + reader = h.endpoint(["A"], slot=4, fill=0x2D) + (result,) = h.settle(h.get(reader.rows())) + assert result.status is RowStatus.FAULT, result + assert reader.intact() # nothing lands outside the row + + def test_an_object_smaller_than_its_row_is_a_fault(self) -> None: + """A store may answer this with an error that leaves access unprovable; the harness + accepts that and retains the memory.""" + with self.harness() as h: + h.needs(put=True, get=True) + h.publish({"A": b"aaaa"}, slot=4) + reader = h.endpoint(["A"], slot=8, fill=0x2D) + (result,) = h.settle(h.get(reader.rows())) + assert result.status is RowStatus.FAULT, result + assert reader.intact() + + def test_a_duplicate_put_of_a_committed_object_is_ready(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True) + h.publish({"A": b"aaaaaaaa"}) + h.publish({"A": b"aaaaaaaa"}) # a second writer, same content + reader = h.endpoint(["A"]) + (result,) = h.settle(h.get(reader.rows())) + assert result.status is RowStatus.READY + assert reader.read_row(0) == b"aaaaaaaa" + + def test_a_put_under_an_existing_key_leaves_the_object_whole(self) -> None: + """Whatever a driver answers a writer bringing other bytes, the object readers get is + the one that was committed.""" + with self.harness() as h: + h.needs(put=True, get=True) + h.publish({"A": b"aaaaaaaa"}) + other = h.endpoint(["A"]) + other.fill_row(0, b"zzzzzzzz") + (result,) = h.settle(h.put(other.rows())) + assert result.status in (RowStatus.READY, RowStatus.FAULT) + reader = h.endpoint(["A"], fill=0x2D) + (got,) = h.settle(h.get(reader.rows())) + assert got.status is RowStatus.READY + assert reader.read_row(0) == b"aaaaaaaa" + + def test_an_operation_leaves_the_rest_of_its_region_alone(self) -> None: + """A row is the only part of its region an operation may touch; holding back only a + stuck operation's own slots depends on it.""" + with self.harness() as h: + h.needs(put=True, get=True) + h.publish({"B": b"bbbbbbbb"}) + reader = h.endpoint(["A", "B", "C"], fill=0x2D) + op = h.get([reader.rows()[1]]) + (result,) = h.settle(op) + assert result.status is RowStatus.READY + h.settle_access(op) + assert reader.read_row(1) == b"bbbbbbbb" + assert reader.read_row(0) == b"-" * SLOT and reader.read_row(2) == b"-" * SLOT + assert reader.intact() + + def test_a_multi_span_row_is_stored_and_fetched_back_to_back(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True, multi_span=True) + writer = h.endpoint(["A", "B", "C"], slot=4) + writer.fill_row(0, b"aaaa") + writer.fill_row(1, b"----") + writer.fill_row(2, b"cccc") + row = RowIO(h.key("M"), (writer.span(0), writer.span(2))) # 8 bytes, slots 0 and 2 + (put,) = h.settle(h.put([row])) + assert put.status is RowStatus.READY + + whole = h.endpoint(["M"], slot=8, fill=0x2D) + (got,) = h.settle(h.get(whole.rows())) + assert got.status is RowStatus.READY + assert whole.read_row(0) == b"aaaacccc" + + parts = h.endpoint(["X", "Y", "Z"], slot=4, fill=0x2D) + back = RowIO(h.key("M"), (parts.span(2), parts.span(0))) # given order, not address + (got,) = h.settle(h.get([back])) + assert got.status is RowStatus.READY + assert parts.read_row(2) == b"aaaa" and parts.read_row(0) == b"cccc" + assert parts.read_row(1) == b"----" + assert writer.intact() and whole.intact() and parts.intact() + + def test_a_row_may_cover_parts_of_slots_across_regions(self) -> None: + """Spans need not be whole slots or share a region, and put and get may split a row + differently; only the bytes in the spans move.""" + with self.harness() as h: + h.needs(put=True, get=True, multi_span=True) + left = h.endpoint(["L"], slot=8, fill=0x2D) + right = h.endpoint(["R"], slot=8, fill=0x2D) + left.fill_row(0, b"-xyz----") + right.fill_row(0, b"--UVWXY-") + row = RowIO(h.key("P"), (left.span_at(1, 3), right.span_at(2, 5))) # "xyz" + "UVWXY" + (put,) = h.settle(h.put([row])) + assert put.status is RowStatus.READY + + whole = h.endpoint(["P"], slot=8, fill=0x2D) + (got,) = h.settle(h.get(whole.rows())) + assert got.status is RowStatus.READY + assert whole.read_row(0) == b"xyzUVWXY" + + split = h.endpoint(["S"], slot=16, fill=0x2D) + back = RowIO(h.key("P"), (split.span_at(10, 5), split.span_at(3, 3))) # 5 then 3 + (got,) = h.settle(h.get([back])) + assert got.status is RowStatus.READY + assert split.read_row(0) == b"---WXY----xyzUV-" # the first 5 bytes land first + assert left.intact() and right.intact() and whole.intact() and split.intact() + + # Rejections leave no trace. + + def test_rows_outside_a_live_registration_are_rejected_before_any_effect(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True, registers=True) + inside = h.endpoint(["A"]) + stray = HostBuffer(SLOT, fill=0x2D) + h.publish({"A": b"aaaaaaaa"}) + + with pytest.raises(DriverRejected): + h.get([RowIO(h.key("A"), (stray.span(),))]) + assert stray.read() == b"-" * SLOT and stray.canaries_intact() + + inside.fill_row(0, b"bbbbbbbb") + inside.unregister() + with pytest.raises(DriverRejected): + h.put([RowIO(h.key("B"), (inside.span(0),))]) + reader = h.endpoint(["B"], fill=0x2D) + (result,) = h.settle(h.get(reader.rows())) + h.assert_absent(result, reader, 0) + + def test_a_key_appears_at_most_once_per_operation(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True, rows=2) + writer = h.endpoint(["A", "A"]) + writer.fill_row(0, b"aaaaaaaa") + writer.fill_row(1, b"bbbbbbbb") + with pytest.raises(DriverRejected): + h.put(writer.rows()) + reader = h.endpoint(["A", "A"], fill=0x2D) + with pytest.raises(DriverRejected): + h.get(reader.rows()) + single = h.endpoint(["A"], fill=0x2D) + (result,) = h.settle(h.get(single.rows())) + h.assert_absent(result, single, 0) + assert reader.read_row(0) == reader.read_row(1) == b"-" * SLOT + + def test_rows_beyond_the_declared_limit_are_rejected(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True) + limit = h.driver.capabilities.max_rows_per_op + if limit is None or limit >= 3: + pytest.skip("this driver declares no row limit this check can exceed") + writer = h.endpoint(["A", "B", "C"][: limit + 1]) + with pytest.raises(DriverRejected): + h.put(writer.rows()) + reader = h.endpoint(["A"], fill=0x2D) + (result,) = h.settle(h.get(reader.rows())) + h.assert_absent(result, reader, 0) + + def test_an_undeclared_direction_is_rejected(self) -> None: + with self.harness() as h: + caps = h.driver.capabilities + if caps.get and caps.put: + pytest.skip("this driver declares both directions") + endpoint = h.endpoint(["A"]) + endpoint.fill_row(0, b"aaaaaaaa") + with pytest.raises(DriverRejected): + if caps.put: + h.get(endpoint.rows()) + else: + h.put(endpoint.rows()) + assert endpoint.read_row(0) == b"aaaaaaaa" + + # Results, access and capabilities over time. + + def test_results_are_stable_and_access_ends_with_evidence(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True, rows=2) + h.publish({"A": b"aaaaaaaa"}) + reader = h.endpoint(["A", "C"]) + op = h.get(reader.rows()) + results = h.settle(op) + assert h.driver.poll(op) == results + h.settle_access(op) + assert h.driver.access_ended(op) + assert not h.driver.access_unprovable(op) + assert h.driver.poll(op) == results + + def test_cancel_after_conclusion_changes_nothing(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True, cancel=True) + h.publish({"A": b"aaaaaaaa"}) + reader = h.endpoint(["A"], fill=0x2D) + op = h.get(reader.rows()) + results = h.settle(op) + h.driver.cancel(op) + assert h.driver.poll(op) == results + assert reader.read_row(0) == b"aaaaaaaa" + + def test_capabilities_do_not_change(self) -> None: + with self.harness() as h: + h.needs(put=True) + before = h.driver.capabilities + h.publish({"A": b"aaaaaaaa"}) + assert h.driver.capabilities == before + assert h.driver.capabilities == before + + def test_a_raise_after_the_operation_is_recorded_becomes_a_row_fault(self) -> None: + driver = self.make_crash_after_record() + if driver is None: + pytest.skip("this driver cannot stage a store that raises mid-operation") + with self.harness(driver) as h: + h.needs(put=True) + writer = h.endpoint(["A"]) + writer.fill_row(0, b"aaaaaaaa") + try: + op = h.put(writer.rows()) + except Exception as error: # noqa: BLE001 - the check is that nothing escapes + pytest.fail(f"put raised after the operation was recorded: {error!r}") + (result,) = h.settle(op) + assert result.status is RowStatus.FAULT, result + + def test_a_key_held_by_an_unfinished_write_is_not_ready_for_get_put_and_probe(self) -> None: + staged = self.make_held_key() + if staged is None: + pytest.skip("this driver cannot stage a key another writer holds") + driver, key, release = staged + with self.harness(driver) as h: + h.needs(put=True, get=True) + early = h.endpoint(["H"], fill=0x2D) + writer = h.endpoint(["H"]) + writer.fill_row(0, b"hhhhhhhh") + write = [RowIO(key, (writer.span(0),))] + try: + (got,) = h.settle(h.get([RowIO(key, (early.span(0),))])) + assert got.status is RowStatus.NOT_READY, got + assert early.read_row(0) == b"-" * SLOT + (put,) = h.settle(h.put(write)) + assert put.status is RowStatus.NOT_READY, put + if driver.capabilities.probe: + assert driver.probe([key], h.timeout) == (ProbeStatus.NOT_READY,) + finally: + release() + (put,) = h.settle(h.put(write)) + assert put.status is RowStatus.READY + # Fresh memory: the spans of the early get stay its own until its access provably ended. + later = h.endpoint(["H"], fill=0x2D) + (got,) = h.settle(h.get([RowIO(key, (later.span(0),))])) + assert got.status is RowStatus.READY + assert later.read_row(0) == b"hhhhhhhh" + + def test_a_row_with_no_room_is_declined_while_the_others_go_on(self) -> None: + driver = self.make_full_store() + if driver is None: + pytest.skip("this driver cannot stage a store with room for one more object") + with self.harness(driver) as h: + h.needs(put=True, get=True, rows=2) + writer = h.endpoint(["A", "B"]) + writer.fill_row(0, b"aaaaaaaa") + writer.fill_row(1, b"bbbbbbbb") + op = h.put(writer.rows()) + results = h.settle(op) + assert sorted((r.status for r in results), key=lambda s: s.value) == [ + RowStatus.DECLINED, + RowStatus.READY, + ], results + h.settle_access(op) + stored = [r.status for r in results].index(RowStatus.READY) + reader = h.endpoint(["A", "B"], fill=0x2D) + got = h.settle(h.get(reader.rows())) + assert got[stored].status is RowStatus.READY + assert reader.read_row(stored) == writer.read_row(stored) + assert got[1 - stored].status is not RowStatus.READY + + def test_a_blocked_get_delays_neither_another_get_nor_the_non_blocking_calls(self) -> None: + staged = self.make_blocked_get() + if staged is None: + pytest.skip("this driver cannot stage a get blocked in the store") + driver, blocked_key, blocked, unblock = staged + with self.harness(driver) as h: + h.needs(put=True, get=True) + h.publish({"B": b"bbbbbbbb"}) + slow = h.endpoint(["S"], fill=0x2D) + quick = h.endpoint(["B"], fill=0x2D) + seen: dict[str, object] = {} + failures: list[BaseException] = [] + + def slow_get() -> None: + try: + h.get([RowIO(blocked_key, (slow.span(0),))]) + except BaseException as error: # noqa: BLE001 - reported from the main thread + failures.append(error) + + def quick_calls() -> None: + try: + op = h.get(quick.rows()) + seen["quick"] = h.settle(op) + seen["ended"] = h.driver.access_ended(op) + seen["drain"] = h.driver.drain(0.05) + except BaseException as error: # noqa: BLE001 - reported from the main thread + failures.append(error) + + # Both off this thread, so a driver that serializes fails instead of hanging. + slow_thread = threading.Thread(target=slow_get, daemon=True) + quick_thread = threading.Thread(target=quick_calls, daemon=True) + slow_thread.start() + try: + if blocked is not None: + assert blocked.wait(h.timeout), "the slow get never reached the store" + else: + time.sleep(0.05) + quick_thread.start() + quick_thread.join(h.timeout) + assert not quick_thread.is_alive(), ( + "another get or a non-blocking call waited behind a get blocked in the store" + ) + if failures: + raise AssertionError(f"a call raised: {failures[0]!r}") from failures[0] + quick_results = seen["quick"] + assert isinstance(quick_results, tuple) + assert [r.status for r in quick_results] == [RowStatus.READY] + assert quick.read_row(0) == b"bbbbbbbb" + assert seen["drain"] is False # the slow get is still inside the store + finally: + unblock() + slow_thread.join(h.timeout) + if quick_thread.ident is not None: + quick_thread.join(h.timeout) + assert not slow_thread.is_alive() and not quick_thread.is_alive() + if failures: + raise AssertionError(f"a call raised: {failures[0]!r}") from failures[0] + + # Probe. + + def test_probe_reports_hit_and_miss(self) -> None: + with self.harness() as h: + h.needs(put=True, probe=True) + h.publish({"A": b"aaaaaaaa"}) + assert h.driver.probe([h.key("A"), h.key("C")], self.timeout) == ( + ProbeStatus.HIT, + ProbeStatus.MISS, + ) + + def test_probe_returns_within_its_budget(self) -> None: + with self.harness() as h: + h.needs(probe=True) + keys = [h.key(label) for label in "ABC"] + started = time.monotonic() + answer = h.driver.probe(keys, 0.05) + elapsed = time.monotonic() - started + assert elapsed < 0.05 + 0.5, f"probe took {elapsed:.2f}s on a 0.05s budget" + assert len(answer) == 3 and set(answer) <= set(ProbeStatus) + + def test_probe_answers_unknown_for_a_key_the_store_does_not_answer_in_time(self) -> None: + staged = self.make_blocked_probe() + if staged is None: + pytest.skip("this driver cannot stage a blocked query") + driver, blocked_key, unblock = staged + with self.harness(driver) as h: + h.needs(probe=True) + answer: list[tuple[ProbeStatus, ...]] = [] + # Off this thread, so a probe that waits for the store fails instead of hanging. + prober = threading.Thread( + target=lambda: answer.append(driver.probe([blocked_key], 0.05)), daemon=True + ) + prober.start() + prober.join(0.05 + 0.5) + try: + assert not prober.is_alive(), "probe waited for the store past its budget" + assert answer == [(ProbeStatus.UNKNOWN,)] + finally: + unblock() + prober.join(h.timeout) + + def test_probe_raises_on_a_store_error_rather_than_answering(self) -> None: + staged = self.make_failing_presence() + if staged is None: + pytest.skip("this driver cannot stage a failing presence query") + driver, key = staged + with self.harness(driver) as h: + h.needs(probe=True) + try: + answer = driver.probe([key], h.timeout) + except DriverRejected as error: + pytest.fail(f"a store error was reported as a rejection: {error!r}") + except Exception: # noqa: BLE001 - any store error will do + return + pytest.fail(f"probe answered {answer} on a store error") + + # Registration, drain and close. + + def test_unregister_is_idempotent_and_frees_the_range(self) -> None: + with self.harness() as h: + h.needs(registers=True) + endpoint = h.endpoint(["A"]) + with pytest.raises(DriverRejected): + h.driver.register(endpoint.memory.address, SLOT, RegionLayout(SLOT, 1)) + endpoint.unregister() + endpoint.unregister() + try: + handle = h.driver.register(endpoint.memory.address, SLOT, RegionLayout(SLOT, 1)) + except DriverRejected as error: + pytest.fail(f"the range stayed registered: {error}") + h.driver.unregister(handle) + + def test_drain_is_one_way(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True) + writer = h.publish({"A": b"aaaaaaaa"}) + assert h.driver.drain(self.timeout) is True + with pytest.raises(DriverRejected): + h.put(writer.rows()) + with pytest.raises(DriverRejected): + h.get(writer.rows()) + assert h.driver.drain(self.timeout) is True + assert writer.read_row(0) == b"aaaaaaaa" + + def test_close_needs_a_true_drain_and_no_live_registrations(self) -> None: + with self.harness() as h: + h.needs(put=True, registers=True) + writer = h.publish({"A": b"aaaaaaaa"}) + with pytest.raises(DriverRejected): + h.driver.close() + assert h.driver.drain(self.timeout) is True + with pytest.raises(DriverRejected): + h.driver.close() + writer.unregister() + h.driver.close() + h.driver.close() + h.done = True # closed by the check itself + with pytest.raises(DriverRejected): + h.put(writer.rows()) diff --git a/tensorrt_llm/_torch/disaggregation/testing/host_memory.py b/tensorrt_llm/_torch/disaggregation/testing/host_memory.py new file mode 100644 index 000000000000..6541fe798769 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/testing/host_memory.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host memory with real addresses for driver tests, guarded by canaries on both sides.""" + +from __future__ import annotations + +import ctypes +from typing import Protocol, runtime_checkable + +from tensorrt_llm._torch.disaggregation.blob.store import Span + +CANARY = bytes.fromhex("deadbeef") * 4 +GUARD = len(CANARY) + + +@runtime_checkable +class Memory(Protocol): + """What the conformance harness needs of a piece of memory a driver may read and write.""" + + @property + def address(self) -> int: ... + + @property + def nbytes(self) -> int: ... + + def read(self, offset: int = 0, nbytes: int | None = None) -> bytes: ... + + def write(self, data: bytes, offset: int = 0) -> None: ... + + def canaries_intact(self) -> bool: ... + + +class MemoryWindow: + """``nbytes`` of memory at a given address, with a canary written before and after it. + + The caller owns ``GUARD`` bytes on either side of the window. ``owner`` is whatever keeps the + backing memory alive; the window holds a reference to it, so memory a driver may still touch + stays allocated for as long as the window is retained. + """ + + def __init__(self, address: int, nbytes: int, fill: int = 0, *, owner: object = None) -> None: + if nbytes <= 0: + raise ValueError("a memory window needs a positive size") + self._owner = owner + self.address = address + self.nbytes = nbytes + ctypes.memmove(address - GUARD, CANARY, GUARD) + ctypes.memmove(address + nbytes, CANARY, GUARD) + ctypes.memset(address, fill, nbytes) + + def span(self, offset: int = 0, nbytes: int | None = None) -> Span: + return Span(self.address + offset, self.nbytes - offset if nbytes is None else nbytes) + + def read(self, offset: int = 0, nbytes: int | None = None) -> bytes: + return ctypes.string_at( + self.address + offset, self.nbytes - offset if nbytes is None else nbytes + ) + + def write(self, data: bytes, offset: int = 0) -> None: + if offset < 0 or offset + len(data) > self.nbytes: + raise ValueError("write outside the window") + ctypes.memmove(self.address + offset, data, len(data)) + + def canaries_intact(self) -> bool: + before = ctypes.string_at(self.address - GUARD, GUARD) + after = ctypes.string_at(self.address + self.nbytes, GUARD) + return before == CANARY and after == CANARY + + +class HostBuffer(MemoryWindow): + """A freshly allocated window that owns its backing storage for as long as it lives.""" + + def __init__(self, nbytes: int, fill: int = 0) -> None: + if nbytes <= 0: + raise ValueError("a host buffer needs a positive size") + self._raw = ctypes.create_string_buffer(GUARD * 2 + nbytes) + super().__init__(ctypes.addressof(self._raw) + GUARD, nbytes, fill) diff --git a/tests/unittest/_torch/disaggregation/blob/memory_store.py b/tensorrt_llm/_torch/disaggregation/testing/memory_store.py similarity index 84% rename from tests/unittest/_torch/disaggregation/blob/memory_store.py rename to tensorrt_llm/_torch/disaggregation/testing/memory_store.py index f304d699c1db..5b1daac592eb 100644 --- a/tests/unittest/_torch/disaggregation/blob/memory_store.py +++ b/tensorrt_llm/_torch/disaggregation/testing/memory_store.py @@ -3,9 +3,11 @@ """An in-process store and a driver over it, for tests. ``MemoryStore`` is a dictionary of complete objects that endpoints in one process can share. It -can be given a capacity, and a test can mark a key as held by an unfinished write. ``MemoryDriver`` -really copies bytes between registered host memory and that store, keeping a trailer beside each -object that it verifies on every get, the way a real store checks sizes and checksums natively. +can be given a capacity, and a test can mark a key as held by an unfinished write, corrupt an +object or make a presence query fail. ``MemoryDriver`` really copies bytes between registered +host memory and that store, gathering a row's spans into one object and scattering them back, +and keeps a trailer beside each object that it verifies on every get, the way a real store checks +sizes and checksums natively. In automatic mode an operation completes inside ``put``/``get``, like a blocking SDK. In controlled mode the test releases three things separately per operation: the transfer, the logical results @@ -26,15 +28,7 @@ from dataclasses import dataclass from typing import Literal -from object_trailer import Checksum, check_trailer, encode_trailer - from tensorrt_llm._torch.disaggregation.blob.store import ( - CANCELLED, - DECLINED, - MISS, - NOT_READY, - PENDING, - READY, DriverCapabilities, DriverRejected, Hint, @@ -51,13 +45,28 @@ faulted, ) +from .object_trailer import Checksum, check_trailer, encode_trailer + DEFAULT_CAPABILITIES = DriverCapabilities( - get=True, put=True, cancel=True, probe=True, max_rows_per_op=None + get=True, + put=True, + cancel=True, + probe=True, + max_rows_per_op=None, + absent_is_miss=True, + multi_span=True, ) MAX_PROBES_IN_FLIGHT = 4 Kind = Literal["put", "get"] +_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) + class ObjectNotReady(StoreError): """The key is held by a write that has not finished.""" @@ -84,14 +93,16 @@ class StoredObject: class MemoryStore: """Objects by key. Each object is a private copy, committed whole or not at all. - Subclass and override ``commit`` or ``get`` to make a test's store block or fail; both may - raise ``StoreError``, which the driver turns into a row fault. + Subclass and override ``commit``, ``get`` or ``presence`` to make a test's store block or + fail; they may raise ``StoreError``, which the driver turns into a row fault or re-raises from + ``probe``. """ def __init__(self, capacity: int | None = None) -> None: self._lock = threading.Lock() self._objects: dict[str, StoredObject] = {} self._writing: set[str] = set() + self._presence_failure: str | None = None self.capacity = capacity def commit(self, key: str, payload: bytes, trailer: bytes) -> Commit: @@ -115,6 +126,9 @@ def get(self, key: str) -> StoredObject | None: def presence(self, key: str) -> ProbeStatus: with self._lock: + if self._presence_failure is not None: + reason, self._presence_failure = self._presence_failure, None + raise StoreError(reason) if key in self._writing: return ProbeStatus.NOT_READY return ProbeStatus.HIT if key in self._objects else ProbeStatus.MISS @@ -147,6 +161,11 @@ def end_write(self, key: str) -> None: with self._lock: self._writing.discard(key) + def fail_next_presence(self, reason: str) -> None: + """Make the next presence query raise ``StoreError(reason)``.""" + with self._lock: + self._presence_failure = reason + def corrupt(self, key: str, index: int) -> None: """Flip one payload byte of the committed object, as a damaged store would.""" with self._lock: @@ -171,7 +190,7 @@ class _Registration: @dataclass(eq=False) class _Row: io: RowIO - published: RowResult = PENDING + published: RowResult = _PENDING staged: RowResult | None = None # decided but not yet visible through poll started: bool = False # a transfer took the row; cancel and fail leave it alone transferred: bool = False @@ -194,6 +213,17 @@ class _Probe: finished: threading.Event +def _gather(spans: Sequence[Span]) -> bytes: + return b"".join(ctypes.string_at(span.address, span.nbytes) for span in spans) + + +def _scatter(spans: Sequence[Span], payload: bytes) -> None: + offset = 0 + for span in spans: + ctypes.memmove(span.address, payload[offset : offset + span.nbytes], span.nbytes) + offset += span.nbytes + + class MemoryDriver: """A ``StoreDriver`` over a ``MemoryStore``; see the module docstring for the two modes.""" @@ -211,6 +241,7 @@ def __init__( self._registrations: list[_Registration] = [] self._ops: list[MemoryOp] = [] self._probes: list[_Probe] = [] + self._probe_workers: list[threading.Thread] = [] # every worker ever started self._draining = False self._closed = False self._cancel_effective = True @@ -310,10 +341,11 @@ def ask() -> None: # start happen under the lock, so close and drain see every worker. with self._lock: self._probes = [probe for probe in self._probes if not probe.finished.is_set()] - if self._draining or self._probe_unknown or len(self._probes) >= MAX_PROBES_IN_FLIGHT: + if self._draining or self._probe_unknown or not self._probe_slot_free(): return tuple(ProbeStatus.UNKNOWN for _ in keys) worker = threading.Thread(target=ask, name="memory-driver-probe", daemon=True) self._probes.append(_Probe(worker, finished)) + self._probe_workers.append(worker) worker.start() if finished.wait(budget_s): worker.join() # it is exiting; leave no thread behind a call that got its answer @@ -390,7 +422,7 @@ def fail(self, op: MemoryOp, reason: str, rows: Iterable[int] | None = None) -> row.published = faulted(reason) def write_late(self, op: MemoryOp, index: int, data: bytes) -> None: - """Write into a faulted get row's payload after its result, as a transfer that kept going + """Write into a faulted get row's spans after its result, as a transfer that kept going would. Refused once access has ended, and for rows that are not FAULT: those must stay as reported.""" with self._lock: @@ -400,12 +432,9 @@ def write_late(self, op: MemoryOp, index: int, data: bytes) -> None: row = op.rows[index] if row.published.status is not RowStatus.FAULT: raise RuntimeError(f"row {index} is {row.published.status.value}, not faulted") - payload = row.io.payload - if len(data) != payload.nbytes: - raise ValueError( - f"late write of {len(data)} bytes into a {payload.nbytes}-byte row" - ) - ctypes.memmove(payload.address, data, len(data)) + if len(data) != row.io.nbytes: + raise ValueError(f"late write of {len(data)} bytes into a {row.io.nbytes}-byte row") + _scatter(row.io.spans, data) def end_access(self, op: MemoryOp) -> None: """Declare that the operation will never touch its rows' memory again.""" @@ -450,11 +479,12 @@ def probes_in_flight(self) -> int: return sum(1 for probe in self._probes if not probe.finished.is_set()) def join_probes(self, timeout: float) -> None: - """Wait for every probe worker to exit, for tests that opened a gate they had closed.""" + """Wait for every probe worker ever started to exit, for tests that opened a gate they + had closed.""" with self._lock: - probes = list(self._probes) - for probe in probes: - probe.worker.join(timeout) + workers = list(self._probe_workers) + for worker in workers: + worker.join(timeout) # Submission checks, one rule each. @@ -466,6 +496,10 @@ def _check_not_draining(self) -> None: if self._draining: raise DriverRejected("the driver is draining") + def _check_hint(self, hint: Hint) -> None: + if hint is not None and not self._capabilities.hints: + raise DriverRejected("this driver reads no hints") + def _check_row_limit(self, rows: Sequence[RowIO]) -> None: if not rows: raise DriverRejected("an operation needs at least one row") @@ -478,21 +512,33 @@ def _check_unique_keys(self, rows: Sequence[RowIO]) -> None: if len(set(keys)) != len(keys): raise DriverRejected("a key appears more than once in the operation") - def _check_hint(self, hint: Hint) -> None: - if hint is not None and not self._capabilities.hints: - raise DriverRejected("this driver reads no hints") + def _check_span_count(self, rows: Sequence[RowIO]) -> None: + if not self._capabilities.multi_span and any(len(io.spans) > 1 for io in rows): + raise DriverRejected("this driver takes rows of one span") + + def _check_disjoint(self, kind: Kind, rows: Sequence[RowIO]) -> None: + """The spans of a get never overlap across rows; the caller guarantees it, this driver checks + it anyway to catch a backend mistake in tests.""" + if kind != "get": + return + spans = [span for io in rows for span in io.spans] + for i, first in enumerate(spans): + for second in spans[i + 1 :]: + if first.overlaps(second): + raise DriverRejected(f"spans {first} and {second} overlap within one get") def _check_registered(self, rows: Sequence[RowIO]) -> None: for io in rows: - if not self._registered(io.payload): - raise DriverRejected(f"{io.payload} is outside every live registration") + for span in io.spans: + if not self._registered(span): + raise DriverRejected(f"{span} is outside every live registration") def _check_region_idle(self, registration: _Registration) -> None: busy = sum( 1 for op in self._ops if not op.access_ended - and any(row.io.payload.within(registration.span) for row in op.rows) + and any(span.within(registration.span) for row in op.rows for span in row.io.spans) ) if busy: raise DriverRejected(f"{busy} operation(s) may still access the region") @@ -503,6 +549,9 @@ def _check_closable(self) -> None: if any(r.live for r in self._registrations): raise DriverRejected("close needs every registration unregistered") + def _probe_slot_free(self) -> bool: + return len(self._probes) < MAX_PROBES_IN_FLIGHT + def _start_draining(self) -> None: self._draining = True @@ -525,6 +574,8 @@ def _submit(self, kind: Kind, rows: Sequence[RowIO], hint: Hint = None) -> Memor self._check_hint(hint) self._check_row_limit(rows) self._check_unique_keys(rows) + self._check_span_count(rows) + self._check_disjoint(kind, rows) self._check_registered(rows) op = MemoryOp(kind, [_Row(io) for io in rows], hint) self._ops.append(op) @@ -565,15 +616,15 @@ def _move(self, kind: Kind, io: RowIO) -> RowResult: return faulted(f"{type(error).__name__}: {error}") def _put_row(self, io: RowIO) -> RowResult: - payload = ctypes.string_at(io.payload.address, io.payload.nbytes) + payload = _gather(io.spans) trailer = self._trailer_for(io, payload) outcome = self.store.commit(io.key, payload, trailer) if outcome in (Commit.STORED, Commit.SAME): - return READY + return _READY if outcome is Commit.DECLINED: - return DECLINED + return _DECLINED if outcome is Commit.NOT_READY: - return NOT_READY + return _NOT_READY # A diagnostic this store can afford; the contract asks no driver to compare bytes. return faulted("key exists with different content") @@ -581,30 +632,28 @@ def _get_row(self, io: RowIO) -> RowResult: try: stored = self.store.get(io.key) except ObjectNotReady: - return NOT_READY + return _NOT_READY if stored is None: - return MISS - if len(stored.payload) != io.payload.nbytes: - return faulted(f"object is {len(stored.payload)} bytes, expected {io.payload.nbytes}") + return _MISS + if len(stored.payload) != io.nbytes: + return faulted(f"object is {len(stored.payload)} bytes, expected {io.nbytes}") problem = self._verify(io, stored) if problem is not None: return faulted(problem) - ctypes.memmove(io.payload.address, stored.payload, len(stored.payload)) - return READY + _scatter(io.spans, stored.payload) + return _READY def _trailer_for(self, io: RowIO, payload: bytes) -> bytes: - return encode_trailer( - io.key.encode(), io.payload.nbytes, checksum=Checksum.CRC32, payload=payload - ) + return encode_trailer(io.key.encode(), io.nbytes, checksum=Checksum.CRC32, payload=payload) def _verify(self, io: RowIO, stored: StoredObject) -> str | None: """Why the stored object may not be delivered as ``io``, or ``None``.""" - return check_trailer(stored.trailer, io.key.encode(), io.payload.nbytes, stored.payload) + return check_trailer(stored.trailer, io.key.encode(), io.nbytes, stored.payload) def _cancel_unstarted(self, op: MemoryOp) -> None: for row in op.rows: if not row.started and not row.published.concluded: - row.published = CANCELLED + row.published = _CANCELLED self._settle_untouched(op) def _settle_untouched(self, op: MemoryOp) -> None: diff --git a/tests/unittest/_torch/disaggregation/blob/object_trailer.py b/tensorrt_llm/_torch/disaggregation/testing/object_trailer.py similarity index 100% rename from tests/unittest/_torch/disaggregation/blob/object_trailer.py rename to tensorrt_llm/_torch/disaggregation/testing/object_trailer.py diff --git a/tests/unittest/_torch/disaggregation/blob/driver_conformance.py b/tests/unittest/_torch/disaggregation/blob/driver_conformance.py deleted file mode 100644 index 4faafe92a11e..000000000000 --- a/tests/unittest/_torch/disaggregation/blob/driver_conformance.py +++ /dev/null @@ -1,505 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""What any ``StoreDriver`` owes the backend, checked through the protocol alone. - -A driver's test module subclasses :class:`DriverConformance` as ``Test...`` and provides -``make_driver``. Every check observes consequences the backend could observe: bytes in registered -memory, row results, access end, rejections without effect. Nothing here uses a driver's own -controls, so the same checks run against the in-memory driver, a real store's driver, and the -deliberately broken drivers that prove each check bites. - -Each check runs inside a :class:`Harness` that keeps every endpoint and operation it created and, -when the check is over, waits for access to end, unregisters the regions, drains and closes the -driver. Memory whose access never provably ended is kept alive instead of being freed. -""" - -from __future__ import annotations - -import contextlib -import hashlib -import threading -import time -from collections.abc import Callable, Generator, Sequence - -import pytest -from host_memory import HostBuffer - -from tensorrt_llm._torch.disaggregation.blob.keys import encode_key -from tensorrt_llm._torch.disaggregation.blob.store import ( - DriverRejected, - Hint, - OpHandle, - ProbeStatus, - RegionHandle, - RegionLayout, - RegistrationError, - RowIO, - RowResult, - RowStatus, - Span, - StoreDriver, - concluded, -) - -SLOT = 8 -PREFIX = "t" - -# Lender names are 54 bytes and start with a 16-byte scope digest shared by every name of a -# deployment. A key built from a truncated name collides across these. -SCOPE = b"scope" + bytes(11) -NAME_BYTES = 54 - - -def name_of(label: str) -> bytes: - return SCOPE + hashlib.blake2b(label.encode(), digest_size=NAME_BYTES - len(SCOPE)).digest() - - -NAMES = {label: name_of(label) for label in "ABC"} - - -def key_of(label: str) -> str: - return encode_key(PREFIX, NAMES[label]) - - -def wait_until(predicate: Callable[[], bool], what: str, timeout: float) -> None: - """Polls ``predicate`` until true; fails after ``timeout`` seconds.""" - deadline = time.monotonic() + timeout - while not predicate(): - assert time.monotonic() < deadline, f"timed out waiting for {what}" - time.sleep(0.001) - - -def statuses(driver: StoreDriver, op: OpHandle) -> list[RowStatus]: - return [result.status for result in driver.poll(op)] - - -class Endpoint: - """One side of a transfer: one registered region with a slot per label.""" - - def __init__( - self, - driver: StoreDriver, - labels: Sequence[str], - *, - fill: int = 0, - slot: int = SLOT, - register: bool = True, - ) -> None: - self.driver = driver - self.labels = list(labels) - self.slot = slot - self.payload = HostBuffer(slot * len(labels), fill) - self.handle: RegionHandle | None = None - if register: - self.register() - - def register(self) -> None: - """Registers the region. When the store may have kept it despite an error, the handle - is recorded all the same, so the region can still be unregistered later.""" - layout = RegionLayout(self.slot, len(self.labels), "kv") - try: - self.handle = self.driver.register(self.payload.address, self.payload.nbytes, layout) - except RegistrationError as error: - self.handle = error.handle - raise - - def rows(self) -> list[RowIO]: - return [ - RowIO(key_of(label), self.span(i), position=(0, i)) - for i, label in enumerate(self.labels) - ] - - def span(self, index: int) -> Span: - return self.payload.span(index * self.slot, self.slot) - - def fill_row(self, index: int, data: bytes) -> None: - self.payload.write(data, index * self.slot) - - def read_row(self, index: int) -> bytes: - return self.payload.read(index * self.slot, self.slot) - - def intact(self) -> bool: - return self.payload.canaries_intact() - - def unregister(self) -> None: - if self.handle is not None: - self.driver.unregister(self.handle) - - -# Endpoints whose memory a driver may still touch or still holds registered when a check ends. -# Never freed. -UNRELEASED_ENDPOINTS: list[Endpoint] = [] - - -class Harness: - """One driver under test plus everything a check created with it.""" - - def __init__(self, driver: StoreDriver, timeout: float) -> None: - self.driver = driver - self.timeout = timeout - self.endpoints: list[Endpoint] = [] - self.ops: list[OpHandle] = [] - self.done = False - - # Creating things the harness will clean up. - - def endpoint(self, labels: Sequence[str], *, fill: int = 0, slot: int = SLOT) -> Endpoint: - endpoint = Endpoint(self.driver, labels, fill=fill, slot=slot, register=False) - self.endpoints.append(endpoint) # kept before the store is asked, whatever register does - endpoint.register() - return endpoint - - def put(self, rows: Sequence[RowIO]) -> OpHandle: - op = self.driver.put(rows) - self.ops.append(op) - return op - - def get(self, rows: Sequence[RowIO], hint: Hint = None) -> OpHandle: - # Without a hint the call keeps its old shape, as the backend's does. - op = self.driver.get(rows) if hint is None else self.driver.get(rows, hint) - self.ops.append(op) - return op - - def publish(self, rows: dict[str, bytes], slot: int = SLOT) -> Endpoint: - """Puts ``rows`` (label to payload) and returns the writer endpoint.""" - writer = self.endpoint(list(rows), slot=slot) - for i, data in enumerate(rows.values()): - writer.fill_row(i, data) - op = self.put(writer.rows()) - assert [r.status for r in self.settle(op)] == [RowStatus.READY] * len(rows) - self.settle_access(op) - return writer - - # Waiting. - - def settle(self, op: OpHandle) -> tuple[RowResult, ...]: - """Waits for every row's result, not for access end.""" - wait_until( - lambda: concluded(self.driver.poll(op)), "the operation to conclude", self.timeout - ) - return self.driver.poll(op) - - def settle_access(self, op: OpHandle) -> None: - wait_until(lambda: self.driver.access_ended(op), "access to end", self.timeout) - - def needs(self, *, put: bool = False, get: bool = False, rows: int = 1) -> None: - """Skips the check when the driver does not declare what it needs.""" - caps = self.driver.capabilities - if put and not caps.put: - pytest.skip("this driver declares no put") - if get and not caps.get: - pytest.skip("this driver declares no get") - if caps.max_rows_per_op is not None and rows > caps.max_rows_per_op: - pytest.skip(f"this driver takes at most {caps.max_rows_per_op} rows per operation") - - # Cleanup. - - def cleanup(self) -> list[str]: - """Ends everything the check created, in order; returns what did not go to plan. - - Memory a driver may still touch, or still holds registered, is retained for good. A - cleanup that did not finish may be called again. - """ - if self.done: - return [] - problems: list[str] = [] - for op in self.ops: - try: - self.settle_access(op) - except Exception as error: # noqa: BLE001 - every cleanup step is reported - problems.append(f"access end: {error!r}") - if problems: - UNRELEASED_ENDPOINTS.extend(self.endpoints) - return problems - for endpoint in self.endpoints: - try: - endpoint.unregister() - except Exception as error: # noqa: BLE001 - every cleanup step is reported - problems.append(f"unregister: {error!r}") - UNRELEASED_ENDPOINTS.append(endpoint) - if problems: - return problems - try: - if self.driver.drain(self.timeout): - self.driver.close() - else: - problems.append("drain did not return True") - except Exception as error: # noqa: BLE001 - every cleanup step is reported - problems.append(f"drain/close: {error!r}") - self.done = not problems - return problems - - -class DriverConformance: - """Contract checks for one driver. Subclass as ``Test...`` and implement ``make_driver``.""" - - timeout = 5.0 - - def make_driver(self) -> StoreDriver: - raise NotImplementedError - - def make_blocked_probe(self) -> tuple[StoreDriver, str, Callable[[], None]] | None: - """A driver whose store stalls every query of the returned key until the returned - callable runs; ``None`` when this driver cannot stage that, which skips the check.""" - return None - - @contextlib.contextmanager - def harness(self, driver: StoreDriver | None = None) -> Generator[Harness, None, None]: - """A fresh driver (or ``driver``) with cleanup afterwards. A failing check keeps its own - error; cleanup trouble fails only a check that passed.""" - h = Harness(driver if driver is not None else self.make_driver(), self.timeout) - try: - yield h - except BaseException: - h.cleanup() - raise - problems = h.cleanup() - assert not problems, "cleanup after the check: " + "; ".join(problems) - - # Bytes and results. - - def test_put_then_get_moves_whole_rows_and_never_aliases_the_source(self) -> None: - with self.harness() as h: - h.needs(put=True, get=True, rows=2) - writer = h.publish({"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) - writer.fill_row(0, b"xxxxxxxx") # the source changes after the put concluded - - reader = h.endpoint(["A", "B"], fill=0x2D) - results = h.settle(h.get(reader.rows())) - assert [r.status for r in results] == [RowStatus.READY, RowStatus.READY] - assert reader.read_row(0) == b"aaaaaaaa" - assert reader.read_row(1) == b"bbbbbbbb" - assert reader.intact() and writer.intact() - - def test_get_reports_hit_hit_miss_and_leaves_the_miss_untouched(self) -> None: - with self.harness() as h: - h.needs(put=True, get=True, rows=3) - h.publish({"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) - reader = h.endpoint(["A", "B", "C"], fill=0x2D) - results = h.settle(h.get(reader.rows())) - assert [r.status for r in results] == [ - RowStatus.READY, - RowStatus.READY, - RowStatus.MISS, - ] - assert reader.read_row(2) == b"-" * SLOT - assert reader.intact() - - def test_an_object_of_the_wrong_size_is_a_fault_not_a_miss(self) -> None: - with self.harness() as h: - h.needs(put=True, get=True) - h.publish({"A": b"aaaa"}, slot=4) - reader = h.endpoint(["A"], fill=0x2D) - (result,) = h.settle(h.get(reader.rows())) - assert result.status is RowStatus.FAULT - assert reader.read_row(0) == b"-" * SLOT - assert reader.intact() - - def test_a_duplicate_put_of_a_committed_object_is_ready(self) -> None: - with self.harness() as h: - h.needs(put=True, get=True) - h.publish({"A": b"aaaaaaaa"}) - h.publish({"A": b"aaaaaaaa"}) # a second writer, same content - reader = h.endpoint(["A"]) - (result,) = h.settle(h.get(reader.rows())) - assert result.status is RowStatus.READY - assert reader.read_row(0) == b"aaaaaaaa" - - def test_a_put_under_an_existing_key_leaves_the_object_whole(self) -> None: - """Whatever a driver answers a writer bringing other bytes, the object readers get is - the one that was committed.""" - with self.harness() as h: - h.needs(put=True, get=True) - h.publish({"A": b"aaaaaaaa"}) - other = h.endpoint(["A"]) - other.fill_row(0, b"zzzzzzzz") - (result,) = h.settle(h.put(other.rows())) - assert result.status in (RowStatus.READY, RowStatus.FAULT) - reader = h.endpoint(["A"], fill=0x2D) - (got,) = h.settle(h.get(reader.rows())) - assert got.status is RowStatus.READY - assert reader.read_row(0) == b"aaaaaaaa" - - # Rejections leave no trace. - - def test_rows_outside_a_live_registration_are_rejected_before_any_effect(self) -> None: - with self.harness() as h: - h.needs(put=True, get=True) - inside = h.endpoint(["A"]) - stray = HostBuffer(SLOT, fill=0x2D) - h.publish({"A": b"aaaaaaaa"}) - - with pytest.raises(DriverRejected): - h.get([RowIO(key_of("A"), stray.span())]) - assert stray.read() == b"-" * SLOT and stray.canaries_intact() - - inside.fill_row(0, b"bbbbbbbb") - inside.unregister() - with pytest.raises(DriverRejected): - h.put([RowIO(key_of("B"), inside.span(0))]) - reader = h.endpoint(["B"]) - assert [r.status for r in h.settle(h.get(reader.rows()))] == [RowStatus.MISS] - - def test_a_key_appears_at_most_once_per_operation(self) -> None: - with self.harness() as h: - h.needs(put=True, get=True, rows=2) - writer = h.endpoint(["A", "A"]) - writer.fill_row(0, b"aaaaaaaa") - writer.fill_row(1, b"bbbbbbbb") - with pytest.raises(DriverRejected): - h.put(writer.rows()) - reader = h.endpoint(["A", "A"], fill=0x2D) - with pytest.raises(DriverRejected): - h.get(reader.rows()) - single = h.endpoint(["A"]) - assert [r.status for r in h.settle(h.get(single.rows()))] == [RowStatus.MISS] - assert reader.read_row(0) == reader.read_row(1) == b"-" * SLOT - - def test_rows_beyond_the_declared_limit_are_rejected(self) -> None: - with self.harness() as h: - h.needs(put=True, get=True) - limit = h.driver.capabilities.max_rows_per_op - if limit is None or limit >= 3: - pytest.skip("this driver declares no row limit this check can exceed") - writer = h.endpoint(["A", "B", "C"][: limit + 1]) - with pytest.raises(DriverRejected): - h.put(writer.rows()) - reader = h.endpoint(["A"]) - assert [r.status for r in h.settle(h.get(reader.rows()))] == [RowStatus.MISS] - - def test_an_undeclared_direction_is_rejected(self) -> None: - with self.harness() as h: - caps = h.driver.capabilities - if caps.get and caps.put: - pytest.skip("this driver declares both directions") - endpoint = h.endpoint(["A"]) - endpoint.fill_row(0, b"aaaaaaaa") - with pytest.raises(DriverRejected): - if caps.put: - h.get(endpoint.rows()) - else: - h.put(endpoint.rows()) - assert endpoint.read_row(0) == b"aaaaaaaa" - - # Results, access and capabilities over time. - - def test_results_are_stable_and_access_ends_with_evidence(self) -> None: - with self.harness() as h: - h.needs(put=True, get=True, rows=2) - h.publish({"A": b"aaaaaaaa"}) - reader = h.endpoint(["A", "C"]) - op = h.get(reader.rows()) - results = h.settle(op) - assert h.driver.poll(op) == results - h.settle_access(op) - assert h.driver.access_ended(op) - assert not h.driver.access_unprovable(op) - assert h.driver.poll(op) == results - - def test_cancel_after_conclusion_changes_nothing(self) -> None: - with self.harness() as h: - h.needs(put=True, get=True) - if not h.driver.capabilities.cancel: - pytest.skip("this driver declares no cancel") - h.publish({"A": b"aaaaaaaa"}) - reader = h.endpoint(["A"], fill=0x2D) - op = h.get(reader.rows()) - results = h.settle(op) - h.driver.cancel(op) - assert h.driver.poll(op) == results - assert reader.read_row(0) == b"aaaaaaaa" - - def test_capabilities_do_not_change(self) -> None: - with self.harness() as h: - h.needs(put=True) - before = h.driver.capabilities - h.publish({"A": b"aaaaaaaa"}) - assert h.driver.capabilities == before - assert h.driver.capabilities == before - - def test_probe_reports_hit_and_miss(self) -> None: - with self.harness() as h: - h.needs(put=True) - if not h.driver.capabilities.probe: - pytest.skip("this driver declares no probe") - h.publish({"A": b"aaaaaaaa"}) - assert h.driver.probe([key_of("A"), key_of("C")], self.timeout) == ( - ProbeStatus.HIT, - ProbeStatus.MISS, - ) - - def test_probe_returns_within_its_budget(self) -> None: - with self.harness() as h: - if not h.driver.capabilities.probe: - pytest.skip("this driver declares no probe") - keys = [key_of(label) for label in "ABC"] - started = time.monotonic() - answer = h.driver.probe(keys, 0.05) - elapsed = time.monotonic() - started - assert elapsed < 0.05 + 0.5, f"probe took {elapsed:.2f}s on a 0.05s budget" - assert len(answer) == 3 and set(answer) <= set(ProbeStatus) - - def test_probe_answers_unknown_for_a_key_the_store_does_not_answer_in_time(self) -> None: - staged = self.make_blocked_probe() - if staged is None: - pytest.skip("this driver cannot stage a blocked query") - driver, blocked_key, unblock = staged - with self.harness(driver) as h: - if not driver.capabilities.probe: - pytest.skip("this driver declares no probe") - answer: list[tuple[ProbeStatus, ...]] = [] - # Off this thread, so a probe that waits for the store fails instead of hanging. - prober = threading.Thread( - target=lambda: answer.append(driver.probe([blocked_key], 0.05)), daemon=True - ) - prober.start() - prober.join(0.05 + 0.5) - try: - assert not prober.is_alive(), "probe waited for the store past its budget" - assert answer == [(ProbeStatus.UNKNOWN,)] - finally: - unblock() - prober.join(h.timeout) - - # Registration, drain and close. - - def test_unregister_is_idempotent_and_frees_the_range(self) -> None: - with self.harness() as h: - endpoint = h.endpoint(["A"]) - with pytest.raises(DriverRejected): - h.driver.register(endpoint.payload.address, SLOT, RegionLayout(SLOT, 1)) - endpoint.unregister() - endpoint.unregister() - try: - handle = h.driver.register(endpoint.payload.address, SLOT, RegionLayout(SLOT, 1)) - except DriverRejected as error: - pytest.fail(f"the range stayed registered: {error}") - h.driver.unregister(handle) - - def test_drain_is_one_way(self) -> None: - with self.harness() as h: - h.needs(put=True, get=True) - writer = h.publish({"A": b"aaaaaaaa"}) - assert h.driver.drain(self.timeout) is True - with pytest.raises(DriverRejected): - h.put(writer.rows()) - with pytest.raises(DriverRejected): - h.get(writer.rows()) - assert h.driver.drain(self.timeout) is True - assert writer.read_row(0) == b"aaaaaaaa" - - def test_close_needs_a_true_drain_and_no_live_registrations(self) -> None: - with self.harness() as h: - h.needs(put=True) - writer = h.publish({"A": b"aaaaaaaa"}) - with pytest.raises(DriverRejected): - h.driver.close() - assert h.driver.drain(self.timeout) is True - with pytest.raises(DriverRejected): - h.driver.close() - writer.unregister() - h.driver.close() - h.driver.close() - h.done = True # closed by the check itself - with pytest.raises(DriverRejected): - h.put(writer.rows()) diff --git a/tests/unittest/_torch/disaggregation/blob/host_memory.py b/tests/unittest/_torch/disaggregation/blob/host_memory.py deleted file mode 100644 index 9bfb3fb2534c..000000000000 --- a/tests/unittest/_torch/disaggregation/blob/host_memory.py +++ /dev/null @@ -1,50 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Host buffers with real addresses for driver tests, guarded by canaries on both sides.""" - -from __future__ import annotations - -import ctypes - -from tensorrt_llm._torch.disaggregation.blob.store import Span - -CANARY = bytes.fromhex("deadbeef") * 4 - - -class HostBuffer: - """``nbytes`` of host memory a driver may read and write, with a canary before and after. - - The buffer keeps its own reference to the backing storage, so the address stays valid for as - long as the object lives. - """ - - def __init__(self, nbytes: int, fill: int = 0) -> None: - if nbytes <= 0: - raise ValueError("a host buffer needs a positive size") - guard = len(CANARY) - self._raw = ctypes.create_string_buffer(guard * 2 + nbytes) - base = ctypes.addressof(self._raw) - ctypes.memmove(base, CANARY, guard) - ctypes.memmove(base + guard + nbytes, CANARY, guard) - self.address = base + guard - self.nbytes = nbytes - ctypes.memset(self.address, fill, nbytes) - - def span(self, offset: int = 0, nbytes: int | None = None) -> Span: - return Span(self.address + offset, self.nbytes - offset if nbytes is None else nbytes) - - def read(self, offset: int = 0, nbytes: int | None = None) -> bytes: - return ctypes.string_at( - self.address + offset, self.nbytes - offset if nbytes is None else nbytes - ) - - def write(self, data: bytes, offset: int = 0) -> None: - if offset < 0 or offset + len(data) > self.nbytes: - raise ValueError("write outside the buffer") - ctypes.memmove(self.address + offset, data, len(data)) - - def canaries_intact(self) -> bool: - guard = len(CANARY) - before = ctypes.string_at(self.address - guard, guard) - after = ctypes.string_at(self.address + self.nbytes, guard) - return before == CANARY and after == CANARY diff --git a/tests/unittest/_torch/disaggregation/blob/test_blob_format.py b/tests/unittest/_torch/disaggregation/blob/test_blob_format.py index 0bdb179640be..3c88e0b6e147 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_blob_format.py +++ b/tests/unittest/_torch/disaggregation/blob/test_blob_format.py @@ -5,18 +5,27 @@ from __future__ import annotations import pytest -from driver_conformance import NAMES, SCOPE -from object_trailer import TRAILER_BYTES, TRAILER_VERSION, Checksum, check_trailer, encode_trailer from tensorrt_llm._torch.disaggregation.blob.keys import FORMAT_VERSION, encode_key from tensorrt_llm._torch.disaggregation.blob.store import ( DriverCapabilities, + DriverRejected, + RegionLayout, RowIO, RowResult, RowStatus, Span, + StoreError, faulted, ) +from tensorrt_llm._torch.disaggregation.testing.driver_conformance import NAMES, SCOPE +from tensorrt_llm._torch.disaggregation.testing.object_trailer import ( + TRAILER_BYTES, + TRAILER_VERSION, + Checksum, + check_trailer, + encode_trailer, +) pytestmark = pytest.mark.cpu_only @@ -47,6 +56,12 @@ def test_key_prefix_is_a_store_partition_with_a_narrow_charset() -> None: # Value types. +def test_a_rejection_is_not_a_store_error() -> None: + """Code written for store trouble must never swallow a rejection.""" + assert not issubclass(DriverRejected, StoreError) + assert not issubclass(StoreError, DriverRejected) + + def test_row_result_pairs_fault_text_with_fault() -> None: assert faulted("timeout").fault == "timeout" assert RowResult(RowStatus.READY).concluded and not RowResult(RowStatus.PENDING).concluded @@ -56,18 +71,33 @@ def test_row_result_pairs_fault_text_with_fault() -> None: RowResult(RowStatus.READY, "text") -def test_span_layout_and_row_validate_their_numbers() -> None: +def test_span_and_layout_validate_their_numbers() -> None: with pytest.raises(ValueError): Span(-1, 8) with pytest.raises(ValueError): Span(0, 0) assert Span(0x1000, 64).overlaps(Span(0x1020, 64)) assert Span(0x1010, 16).within(Span(0x1000, 64)) + for slot_bytes, slot_count in ((0, 1), (8, 0), (-8, 1)): + with pytest.raises(ValueError): + RegionLayout(slot_bytes, slot_count) + assert RegionLayout(8, 4).nbytes == 32 + + +def test_a_row_is_a_key_and_disjoint_spans_in_the_order_given() -> None: + spans = [Span(0x2000, 8), Span(0x1000, 8)] + row = RowIO("k", spans) # type: ignore[arg-type] + assert row.spans == (Span(0x2000, 8), Span(0x1000, 8)) # kept in order, not sorted + assert row.nbytes == 16 + with pytest.raises(ValueError): + RowIO("", (Span(0x1000, 8),)) with pytest.raises(ValueError): - RowIO("", Span(0x1000, 8)) + RowIO("k", ()) + with pytest.raises(ValueError): + RowIO("k", (Span(0x1000, 64), Span(0x1020, 64))) -def test_capabilities_are_declared_in_full() -> None: +def test_capabilities_that_grant_are_declared_and_later_ones_default_off() -> None: with pytest.raises(TypeError): DriverCapabilities(get=True, put=True, cancel=False, probe=False) # type: ignore[call-arg] with pytest.raises(ValueError): @@ -75,7 +105,12 @@ def test_capabilities_are_declared_in_full() -> None: with pytest.raises(ValueError): DriverCapabilities(get=True, put=True, cancel=False, probe=False, max_rows_per_op=0) caps = DriverCapabilities(get=True, put=False, cancel=False, probe=False, max_rows_per_op=4) - assert caps.hints is False # a capability added later defaults to not supported + assert (caps.hints, caps.absent_is_miss, caps.device_memory, caps.multi_span) == ( + False, + False, + False, + False, + ) with pytest.raises(AttributeError): caps.put = True # type: ignore[misc] diff --git a/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py index 39194107a188..448e822039cb 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py +++ b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py @@ -1,59 +1,64 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The in-memory driver: the shared conformance checks, its controlled mode, and proof that every -check fails against a driver that breaks exactly the rule it guards.""" +"""The in-memory driver: the shared conformance checks run four ways, its controlled mode, and +proof that every check fails against a driver that breaks exactly the rule it guards.""" from __future__ import annotations -import ctypes +import gc import threading import time +import weakref from collections.abc import Callable, Sequence import pytest -from driver_conformance import ( - SLOT, - UNRELEASED_ENDPOINTS, - DriverConformance, - Endpoint, - Harness, - key_of, - statuses, -) -from host_memory import HostBuffer -from memory_store import ( - DEFAULT_CAPABILITIES, - MAX_PROBES_IN_FLIGHT, - Kind, - MemoryDriver, - MemoryOp, - MemoryStore, - StoredObject, -) -from object_trailer import encode_trailer from tensorrt_llm._torch.disaggregation.blob.store import ( - MISS, DriverCapabilities, DriverRejected, Hint, OpHandle, ProbeStatus, - RegionHandle, RegionLayout, RegistrationError, RowIO, RowResult, RowStatus, + Span, StoreDriver, StoreError, concluded, faulted, ) +from tensorrt_llm._torch.disaggregation.testing.broken_drivers import BROKEN_CONFORMANCE +from tensorrt_llm._torch.disaggregation.testing.driver_conformance import ( + SLOT, + UNRELEASED_ENDPOINTS, + Allocation, + DriverConformance, + Endpoint, + Harness, + Unblock, + key_of, + statuses, +) +from tensorrt_llm._torch.disaggregation.testing.host_memory import GUARD, HostBuffer, MemoryWindow +from tensorrt_llm._torch.disaggregation.testing.memory_store import ( + DEFAULT_CAPABILITIES, + MAX_PROBES_IN_FLIGHT, + Kind, + MemoryDriver, + MemoryOp, + MemoryStore, + StoredObject, +) +from tensorrt_llm._torch.disaggregation.testing.object_trailer import encode_trailer pytestmark = pytest.mark.cpu_only -KEY_A, KEY_B, KEY_C = key_of("A"), key_of("B"), key_of("C") +PREFIX = "m" # the prefix the controlled-mode endpoints below use +KEY_A, KEY_B, KEY_C = (key_of(label, PREFIX) for label in "ABC") +STAGED_KEY = key_of("staged", "stage") TIMEOUT = 5.0 @@ -96,8 +101,8 @@ def commit(self, key: str, payload: bytes, trailer: bytes): # type: ignore[no-u class GatedStore(MemoryStore): - """Holds reads and writes of ``held`` keys (every key when None) while the gate is closed, - and says when one is waiting.""" + """Holds reads, writes and queries of ``held`` keys (every key when None) while the gate is + closed, and says when one is waiting.""" def __init__(self, held: set[str] | None = None) -> None: super().__init__() @@ -130,6 +135,62 @@ def _run(target: Callable[..., object], *args: object) -> threading.Thread: return worker +def endpoint( + driver: StoreDriver, labels: Sequence[str], *, fill: int = 0, slot: int = SLOT +) -> Endpoint: + """A registered endpoint for the controlled-mode tests, outside any harness.""" + made = Endpoint( + driver, labels, HostBuffer(slot * len(labels), fill), None, prefix=PREFIX, slot=slot + ) + made.register() + return made + + +# The conformance suite, four ways. + + +class TestMemoryDriverConformance(DriverConformance): + """The suite plus every scenario the in-memory store can stage.""" + + driver_cls: type[MemoryDriver] = MemoryDriver + capabilities = DEFAULT_CAPABILITIES + + def _driver(self, store: MemoryStore | None = None) -> MemoryDriver: + # An empty store is falsy (it has a length), so test for None explicitly. + return self.driver_cls( + store if store is not None else MemoryStore(), capabilities=self.capabilities + ) + + def make_driver(self) -> StoreDriver: + return self._driver() + + def make_blocked_probe(self) -> tuple[StoreDriver, str, Unblock]: + store = GatedStore(held={STAGED_KEY}) + store.gate.clear() + return self._driver(store), STAGED_KEY, store.gate.set + + def make_blocked_get(self) -> tuple[StoreDriver, str, threading.Event | None, Unblock]: + store = GatedStore(held={STAGED_KEY}) + store.gate.clear() + return self._driver(store), STAGED_KEY, store.waiting, store.gate.set + + def make_held_key(self) -> tuple[StoreDriver, str, Unblock]: + store = MemoryStore() + store.begin_write(STAGED_KEY) + return self._driver(store), STAGED_KEY, lambda: store.end_write(STAGED_KEY) + + def make_full_store(self) -> StoreDriver: + return self._driver(MemoryStore(capacity=1)) + + def make_crash_after_record(self) -> StoreDriver: + return self._driver(CrashingStore()) + + def make_failing_presence(self) -> tuple[StoreDriver, str]: + store = MemoryStore() + store.fail_next_presence("connection dropped") + return self._driver(store), STAGED_KEY + + class _HintFreeSignatures(MemoryDriver): """A driver written before hints existed: ``get`` and ``probe`` take no hint. With ``hints`` undeclared the backend never passes one, so it stays a valid driver.""" @@ -141,21 +202,47 @@ def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...] return super().probe(keys, budget_s) -class TestMemoryDriverConformance(DriverConformance): - driver_cls: type[MemoryDriver] = MemoryDriver - capabilities = DEFAULT_CAPABILITIES +class TestHintFreeDriverConformance(TestMemoryDriverConformance): + driver_cls = _HintFreeSignatures + + +SHARED_STORE = MemoryStore() + + +class TestSharedStoreConformance(DriverConformance): + """Fresh clients of one persistent store: every check finds the previous checks' objects.""" def make_driver(self) -> StoreDriver: - return self.driver_cls(MemoryStore(), capabilities=self.capabilities) + return MemoryDriver(SHARED_STORE) - def make_blocked_probe(self) -> tuple[StoreDriver, str, Callable[[], None]]: - store = GatedStore(held={KEY_A}) - store.gate.clear() - return self.driver_cls(store, capabilities=self.capabilities), KEY_A, store.gate.set +class TestPoolDriverConformance(DriverConformance): + """A driver that registered its memory once, when built; endpoints are carved out of it.""" -class TestHintFreeDriverConformance(TestMemoryDriverConformance): - driver_cls = _HintFreeSignatures + registers_on_demand = False + POOL_BYTES = 1 << 16 + + def pool_driver(self) -> MemoryDriver: + return MemoryDriver(MemoryStore()) + + def make_driver(self) -> StoreDriver: + driver = self.pool_driver() + self.pool = HostBuffer(self.POOL_BYTES) + self.handle = driver.register( + self.pool.address, self.pool.nbytes, RegionLayout(4, self.POOL_BYTES // 4, "pool") + ) + self.next_offset = 0 + return driver + + def allocate(self, nbytes: int, fill: int) -> Allocation: + start = self.next_offset + GUARD + self.next_offset = start + nbytes + GUARD + assert self.next_offset <= self.POOL_BYTES, "the pool is too small for this check" + return MemoryWindow(self.pool.address + start, nbytes, fill, owner=self.pool), self.handle + + def shutdown(self, driver: StoreDriver) -> None: + driver.unregister(self.handle) + driver.close() class TestMemoryDriver: @@ -170,7 +257,7 @@ def make(self, store: MemoryStore | None = None, **kwargs: object) -> MemoryDriv def _publish( self, driver: MemoryDriver, rows: dict[str, bytes], *, keep_registered: bool = False ) -> Endpoint: - writer = Endpoint(driver, list(rows)) + writer = endpoint(driver, list(rows)) for i, data in enumerate(rows.values()): writer.fill_row(i, data) op = driver.put(writer.rows()) @@ -199,17 +286,27 @@ def test_a_corrupted_object_is_a_fault_and_never_lands(self) -> None: driver = self.make(store) self._publish(driver, {"A": b"aaaaaaaa"}) store.corrupt(KEY_A, 3) - reader = Endpoint(driver, ["A"], fill=0x2D) + reader = endpoint(driver, ["A"], fill=0x2D) (result,) = driver.poll(driver.get(reader.rows())) assert result.status is RowStatus.FAULT assert result.fault == "payload checksum mismatch" assert reader.read_row(0) == b"-" * SLOT + def test_an_object_shorter_than_its_trailer_claims_is_a_fault_and_never_lands(self) -> None: + store = MemoryStore() + driver = self.make(store) + store.overwrite(KEY_A, b"aaaa", encode_trailer(KEY_A.encode(), SLOT)) # claims 8 bytes + reader = endpoint(driver, ["A"], fill=0x2D) + (result,) = driver.poll(driver.get(reader.rows())) + assert result.status is RowStatus.FAULT + assert result.fault == f"object is 4 bytes, expected {SLOT}" + assert reader.read_row(0) == b"-" * SLOT + def test_a_different_object_under_the_same_key_is_diagnosed(self) -> None: store = MemoryStore() driver = self.make(store) self._publish(driver, {"A": b"aaaaaaaa"}) - other = Endpoint(driver, ["A"]) + other = endpoint(driver, ["A"]) other.fill_row(0, b"zzzzzzzz") (result,) = driver.poll(driver.put(other.rows())) assert result.status is RowStatus.FAULT @@ -220,11 +317,11 @@ def test_a_retry_that_fails_in_the_store_leaves_the_committed_object_whole(self) store = FailingStore(fail_on_commit=2) driver = self.make(store) self._publish(driver, {"A": b"aaaaaaaa"}) - retry = Endpoint(driver, ["A"]) + retry = endpoint(driver, ["A"]) retry.fill_row(0, b"aaaaaaaa") (result,) = driver.poll(driver.put(retry.rows())) assert result.status is RowStatus.FAULT and result.fault == "connection dropped" - reader = Endpoint(driver, ["A"], fill=0x2D) + reader = endpoint(driver, ["A"], fill=0x2D) (got,) = driver.poll(driver.get(reader.rows())) assert got.status is RowStatus.READY assert reader.read_row(0) == b"aaaaaaaa" @@ -254,43 +351,13 @@ def read() -> None: assert seen_partial == [] assert len(store) == len(keys) - # Statuses that are not faults. - - def test_a_key_held_by_an_unfinished_write_is_not_ready_for_get_put_and_probe(self) -> None: - store = MemoryStore() - driver = self.make(store) - store.begin_write(KEY_A) - reader = Endpoint(driver, ["A", "C"], fill=0x2D) - results = driver.poll(driver.get(reader.rows())) - assert [r.status for r in results] == [RowStatus.NOT_READY, RowStatus.MISS] - assert reader.read_row(0) == b"-" * SLOT - writer = Endpoint(driver, ["A"]) - writer.fill_row(0, b"aaaaaaaa") - assert statuses(driver, driver.put(writer.rows())) == [RowStatus.NOT_READY] - assert driver.probe([KEY_A, KEY_C], TIMEOUT) == (ProbeStatus.NOT_READY, ProbeStatus.MISS) - store.end_write(KEY_A) - assert statuses(driver, driver.put(writer.rows())) == [RowStatus.READY] - assert statuses(driver, driver.get(reader.rows())) == [RowStatus.READY, RowStatus.MISS] - assert reader.read_row(0) == b"aaaaaaaa" - - def test_a_row_the_store_has_no_room_for_is_declined_not_faulted(self) -> None: - driver = self.make(MemoryStore(capacity=1)) - writer = Endpoint(driver, ["A", "B"]) - writer.fill_row(0, b"aaaaaaaa") - writer.fill_row(1, b"bbbbbbbb") - op = driver.put(writer.rows()) - assert statuses(driver, op) == [RowStatus.READY, RowStatus.DECLINED] - assert driver.access_ended(op) - reader = Endpoint(driver, ["A", "B"], fill=0x2D) - assert statuses(driver, driver.get(reader.rows())) == [RowStatus.READY, RowStatus.MISS] - # Result and access end are independent. def test_result_can_come_before_access_end(self) -> None: store = MemoryStore() self._publish(self.make(store), {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) driver = self.make(store, controlled=True) - reader = Endpoint(driver, ["A", "B"], fill=0x2D) + reader = endpoint(driver, ["A", "B"], fill=0x2D) op = driver.get(reader.rows()) assert statuses(driver, op) == [RowStatus.PENDING, RowStatus.PENDING] @@ -313,7 +380,7 @@ def test_complete_accepts_a_one_shot_iterator(self) -> None: store = MemoryStore() self._publish(self.make(store), {"A": b"aaaaaaaa"}) driver = self.make(store, controlled=True) - reader = Endpoint(driver, ["A"], fill=0x2D) + reader = endpoint(driver, ["A"], fill=0x2D) op = driver.get(reader.rows()) driver.complete(op, iter([0])) assert statuses(driver, op) == [RowStatus.READY] @@ -323,7 +390,7 @@ def test_access_can_end_before_the_result_is_published(self) -> None: store = MemoryStore() self._publish(self.make(store), {"A": b"aaaaaaaa"}) driver = self.make(store, controlled=True) - reader = Endpoint(driver, ["A"], fill=0x2D) + reader = endpoint(driver, ["A"], fill=0x2D) op = driver.get(reader.rows()) driver.transfer(op) driver.end_access(op) @@ -337,7 +404,7 @@ def test_a_late_write_lands_only_while_access_is_open(self) -> None: store = MemoryStore() self._publish(self.make(store), {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) driver = self.make(store, controlled=True) - reader = Endpoint(driver, ["A", "B"], fill=0x2D) + reader = endpoint(driver, ["A", "B"], fill=0x2D) op = driver.get(reader.rows()) driver.complete(op, [0]) driver.fail(op, "timeout", [1]) @@ -362,7 +429,7 @@ def test_a_late_write_lands_only_while_access_is_open(self) -> None: def test_a_cancelled_row_cannot_be_written_late(self) -> None: driver = self.make(controlled=True) - reader = Endpoint(driver, ["A"], fill=0x2D) + reader = endpoint(driver, ["A"], fill=0x2D) op = driver.get(reader.rows()) driver.cancel(op) assert statuses(driver, op) == [RowStatus.CANCELLED] @@ -374,7 +441,7 @@ def test_a_cancelled_row_cannot_be_written_late(self) -> None: def test_a_store_timeout_is_a_fault_not_a_miss(self) -> None: driver = self.make(TimeoutStore(KEY_A)) - reader = Endpoint(driver, ["A", "C"], fill=0x2D) + reader = endpoint(driver, ["A", "C"], fill=0x2D) results = driver.poll(driver.get(reader.rows())) assert [r.status for r in results] == [RowStatus.FAULT, RowStatus.MISS] assert results[0].fault == "timeout" @@ -383,7 +450,7 @@ def test_a_store_timeout_is_a_fault_not_a_miss(self) -> None: def test_a_store_error_mid_operation_becomes_a_row_fault_and_keeps_the_handle(self) -> None: store = FailingStore(fail_on_commit=2) driver = self.make(store) - writer = Endpoint(driver, ["A", "B", "C"]) + writer = endpoint(driver, ["A", "B", "C"]) for i, data in enumerate((b"aaaaaaaa", b"bbbbbbbb", b"cccccccc")): writer.fill_row(i, data) op = driver.put(writer.rows()) @@ -395,30 +462,16 @@ def test_a_store_error_mid_operation_becomes_a_row_fault_and_keeps_the_handle(se assert not store.contains(KEY_B) assert store.committed_payload(KEY_C) == b"cccccccc" - def test_a_raise_after_the_operation_is_recorded_becomes_a_row_fault(self) -> None: - driver = self.make(CrashingStore()) - writer = Endpoint(driver, ["A"]) - writer.fill_row(0, b"aaaaaaaa") - try: - op = driver.put(writer.rows()) - except Exception as error: # noqa: BLE001 - the check is that nothing escapes - pytest.fail(f"put raised after the operation was recorded: {error!r}") - (result,) = driver.poll(op) - assert result.status is RowStatus.FAULT - assert result.fault == "RuntimeError: binding returned garbage" - assert driver.access_ended(op) - assert driver.drain(TIMEOUT) is True - # Concurrency. def test_queries_do_not_wait_on_an_operation_inside_the_store(self) -> None: store = GatedStore() driver = self.make(store, controlled=True) - done = Endpoint(driver, ["A"]) + done = endpoint(driver, ["A"]) finished = driver.get(done.rows()) driver.fail(finished, "timeout") - slow = Endpoint(driver, ["B"]) + slow = endpoint(driver, ["B"]) slow.fill_row(0, b"bbbbbbbb") blocked = driver.put(slow.rows()) store.gate.clear() @@ -467,39 +520,13 @@ def query() -> None: assert statuses(driver, blocked) == [RowStatus.READY] assert store.committed_payload(KEY_B) == b"bbbbbbbb" - def test_a_blocked_get_does_not_delay_another_get(self) -> None: - store = GatedStore(held={KEY_A}) - driver = self.make(store) - self._publish(driver, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) - slow = Endpoint(driver, ["A"], fill=0x2D) - quick = Endpoint(driver, ["B"], fill=0x2D) - store.gate.clear() - outcome: dict[str, list[RowStatus]] = {} - - def fetch(name: str, endpoint: Endpoint) -> None: - outcome[name] = statuses(driver, driver.get(endpoint.rows())) - - slow_worker = _run(fetch, "slow", slow) - try: - assert store.waiting.wait(TIMEOUT), "the slow get never reached the store" - quick_worker = _run(fetch, "quick", quick) - quick_worker.join(TIMEOUT) - assert not quick_worker.is_alive(), "a get waited for another get blocked in the store" - assert outcome["quick"] == [RowStatus.READY] - assert quick.read_row(0) == b"bbbbbbbb" - assert "slow" not in outcome - finally: - store.gate.set() - slow_worker.join(TIMEOUT) - assert outcome["slow"] == [RowStatus.READY] - # Cancel. def test_cancel_concludes_unstarted_rows_and_spares_started_ones(self) -> None: store = MemoryStore() self._publish(self.make(store), {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) driver = self.make(store, controlled=True) - reader = Endpoint(driver, ["A", "B"], fill=0x2D) + reader = endpoint(driver, ["A", "B"], fill=0x2D) op = driver.get(reader.rows()) driver.transfer(op, [0]) driver.cancel(op) @@ -511,7 +538,7 @@ def test_cancel_concludes_unstarted_rows_and_spares_started_ones(self) -> None: def test_cancelling_everything_before_it_started_ends_access(self) -> None: driver = self.make(controlled=True) - reader = Endpoint(driver, ["A", "B"]) + reader = endpoint(driver, ["A", "B"]) op = driver.get(reader.rows()) driver.cancel(op) assert statuses(driver, op) == [RowStatus.CANCELLED, RowStatus.CANCELLED] @@ -520,7 +547,7 @@ def test_cancelling_everything_before_it_started_ends_access(self) -> None: def test_an_ineffective_cancel_changes_nothing(self) -> None: driver = self.make(controlled=True) driver.set_cancel_effective(False) - reader = Endpoint(driver, ["A"]) + reader = endpoint(driver, ["A"]) op = driver.get(reader.rows()) driver.cancel(op) assert statuses(driver, op) == [RowStatus.PENDING] @@ -528,34 +555,23 @@ def test_an_ineffective_cancel_changes_nothing(self) -> None: # Probe. - def test_probe_answers_unknown_when_the_store_does_not_answer_in_time(self) -> None: + def test_a_blocked_probe_holds_drain_until_the_store_answers(self) -> None: store = GatedStore(held={KEY_A}) driver = self.make(store) self._publish(driver, {"B": b"bbbbbbbb"}) - keys = [KEY_B, KEY_A, KEY_C] store.gate.clear() - answer: list[tuple[ProbeStatus, ...]] = [] - elapsed: list[float] = [] - - def probe() -> None: - started = time.monotonic() - answer.append(driver.probe(keys, 0.05)) - elapsed.append(time.monotonic() - started) - - prober = _run(probe) - prober.join(1.0) try: - assert not prober.is_alive(), "probe waited for the store past its budget" - # Keys asked before the blocked one were answered; the rest are unknown, not MISS. - assert answer == [(ProbeStatus.HIT, ProbeStatus.UNKNOWN, ProbeStatus.UNKNOWN)] - assert elapsed[0] < 0.5 + assert driver.probe([KEY_B, KEY_A, KEY_C], 0.05) == ( + ProbeStatus.HIT, + ProbeStatus.UNKNOWN, + ProbeStatus.UNKNOWN, + ) # keys before the blocked one were answered; the rest are unknown, not MISS assert driver.probes_in_flight() == 1 assert driver.drain(0.05) is False # the probe is still inside the store finally: store.gate.set() - prober.join(TIMEOUT) driver.join_probes(TIMEOUT) - assert driver.probe(keys, TIMEOUT) == (ProbeStatus.UNKNOWN,) * 3 # draining now + assert driver.probe([KEY_A], TIMEOUT) == (ProbeStatus.UNKNOWN,) # draining now self._shutdown(driver) def test_probes_in_flight_are_capped(self) -> None: @@ -576,12 +592,20 @@ def test_probes_in_flight_are_capped(self) -> None: self._shutdown(driver) assert driver.probes_in_flight() == 0 + def test_probe_unknown_switch(self) -> None: + driver = self.make() + self._publish(driver, {"A": b"aaaaaaaa"}) + driver.set_probe_unknown(True) + assert driver.probe([KEY_A, KEY_C], TIMEOUT) == (ProbeStatus.UNKNOWN,) * 2 + + # Hints and spans the backend must not hand this driver. + def test_memory_driver_rejects_a_hint_it_cannot_read_before_any_effect(self) -> None: """The backend rejects such hints before dispatch; this driver also refuses them itself, so a backend bug cannot make it read with a hint it silently dropped.""" driver = self.make() self._publish(driver, {"A": b"aaaaaaaa"}) - reader = Endpoint(driver, ["A"], fill=0x2D) + reader = endpoint(driver, ["A"], fill=0x2D) with pytest.raises(DriverRejected): driver.get(reader.rows(), {"source": "worker-7"}) assert reader.read_row(0) == b"-" * SLOT @@ -596,7 +620,7 @@ def test_a_hint_reaches_a_driver_that_declares_hints_unchanged(self) -> None: ) driver = self.make(capabilities=hinted) self._publish(driver, {"A": b"aaaaaaaa"}) - reader = Endpoint(driver, ["A"], fill=0x2D) + reader = endpoint(driver, ["A"], fill=0x2D) hint = {"source": "worker-7"} op = driver.get(reader.rows(), hint) assert statuses(driver, op) == [RowStatus.READY] @@ -607,11 +631,32 @@ def test_a_hint_reaches_a_driver_that_declares_hints_unchanged(self) -> None: assert statuses(driver, driver.get(reader.rows())) == [RowStatus.READY] assert len(driver.hints_seen) == 2 # no hint, nothing recorded - def test_probe_unknown_switch(self) -> None: + def test_memory_driver_rejects_overlapping_spans_within_one_get(self) -> None: + """The caller guarantees disjoint spans; this driver checks anyway, for the same reason + it checks hints.""" driver = self.make() - self._publish(driver, {"A": b"aaaaaaaa"}) - driver.set_probe_unknown(True) - assert driver.probe([KEY_A, KEY_C], TIMEOUT) == (ProbeStatus.UNKNOWN,) * 2 + self._publish(driver, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + reader = endpoint(driver, ["A", "B"], fill=0x2D) + rows = [ + RowIO(KEY_A, (reader.span(0),)), + RowIO(KEY_B, (Span(reader.span(0).address + 4, SLOT),)), # overlaps row A's span + ] + with pytest.raises(DriverRejected): + driver.get(rows) + assert reader.read_row(0) == reader.read_row(1) == b"-" * SLOT + assert statuses(driver, driver.get(reader.rows())) == [RowStatus.READY, RowStatus.READY] + + def test_a_single_span_driver_rejects_a_multi_span_row(self) -> None: + single = DriverCapabilities( + get=True, put=True, cancel=True, probe=True, max_rows_per_op=None, absent_is_miss=True + ) + driver = self.make(capabilities=single) + writer = endpoint(driver, ["A", "B"]) + with pytest.raises(DriverRejected): + driver.put([RowIO(KEY_A, (writer.span(0), writer.span(1)))]) + assert statuses(driver, driver.put(writer.rows())) == [RowStatus.READY, RowStatus.READY] + + # Submission and drain. def test_a_drain_racing_a_submission_leaves_a_cancelled_operation_with_its_handle( self, @@ -621,7 +666,7 @@ def test_a_drain_racing_a_submission_leaves_a_cancelled_operation_with_its_handl "Racing", (self.driver_cls,), {"_before_execute": lambda self, op: self.drain(0.0)} ) driver = racing(store) - writer = Endpoint(driver, ["A"]) + writer = endpoint(driver, ["A"]) writer.fill_row(0, b"aaaaaaaa") try: op = driver.put(writer.rows()) @@ -632,47 +677,10 @@ def test_a_drain_racing_a_submission_leaves_a_cancelled_operation_with_its_handl assert len(store) == 0 assert driver.drain(TIMEOUT) is True - def test_an_object_shorter_than_its_trailer_claims_is_a_fault_and_never_lands(self) -> None: - store = MemoryStore() - driver = self.make(store) - store.overwrite(KEY_A, b"aaaa", encode_trailer(KEY_A.encode(), SLOT)) # claims 8 bytes - reader = Endpoint(driver, ["A"], fill=0x2D) - (result,) = driver.poll(driver.get(reader.rows())) - assert result.status is RowStatus.FAULT - assert result.fault == f"object is 4 bytes, expected {SLOT}" - assert reader.read_row(0) == b"-" * SLOT - - # The conformance harness. - - def test_harness_keeps_a_buffer_whose_registration_escaped(self) -> None: - driver = self.make() - h = Harness(driver, TIMEOUT) - driver.fail_next_register("timed out after the request", escaped=True) - with pytest.raises(RegistrationError): - h.endpoint(["A"]) - assert len(h.endpoints) == 1 and h.endpoints[0].handle is not None - assert h.cleanup() == [] - assert driver.live_registrations() == [] - - def test_harness_retains_an_endpoint_it_could_not_unregister_and_retries(self) -> None: - driver = self.make() - h = Harness(driver, TIMEOUT) - endpoint = h.endpoint(["A"]) - assert endpoint.handle is not None - driver.fail_next_unregister(endpoint.handle, "busy") - problems = h.cleanup() - assert problems and not h.done - assert endpoint in UNRELEASED_ENDPOINTS - assert driver.live_registrations() == [endpoint.handle] - assert h.cleanup() == [] and h.done - assert driver.live_registrations() == [] - - # Registration and shutdown. - def test_unregister_is_refused_while_an_operation_may_access_the_region(self) -> None: driver = self.make(controlled=True) - reader = Endpoint(driver, ["A"]) - other = Endpoint(driver, ["B"]) + reader = endpoint(driver, ["A"]) + other = endpoint(driver, ["B"]) op = driver.get(reader.rows()) with pytest.raises(DriverRejected): reader.unregister() @@ -690,7 +698,7 @@ def test_drain_waits_for_a_started_transfer(self) -> None: store = GatedStore(held={KEY_A}) driver = self.make(store) self._publish(driver, {"A": b"aaaaaaaa"}) - reader = Endpoint(driver, ["A"], fill=0x2D) + reader = endpoint(driver, ["A"], fill=0x2D) store.gate.clear() worker = _run(driver.get, reader.rows()) try: @@ -708,7 +716,7 @@ def test_drain_waits_for_a_started_transfer(self) -> None: def test_an_unprovable_operation_lets_drain_return_false_at_once(self) -> None: driver = self.make(controlled=True) - reader = Endpoint(driver, ["A"]) + reader = endpoint(driver, ["A"]) op = driver.get(reader.rows()) driver.transfer(op) driver.publish_results(op) @@ -733,8 +741,8 @@ def test_drain_returns_false_at_once_even_while_another_operation_is_still_prova self, ) -> None: driver = self.make(controlled=True) - lost = Endpoint(driver, ["A"]) - pending = Endpoint(driver, ["B"]) + lost = endpoint(driver, ["A"]) + pending = endpoint(driver, ["B"]) unprovable = driver.get(lost.rows()) driver.transfer(unprovable) driver.publish_results(unprovable) @@ -750,6 +758,8 @@ def test_drain_returns_false_at_once_even_while_another_operation_is_still_prova driver.end_access(unprovable) assert driver.drain(TIMEOUT) is True + # Registration and shutdown. + def test_register_checks_geometry_and_overlap(self) -> None: driver = self.make() buffer = HostBuffer(64) @@ -800,100 +810,119 @@ def test_calls_after_close_are_programming_errors(self) -> None: with pytest.raises(RuntimeError, match="closed"): driver.probe([KEY_A], 0.0) with pytest.raises(DriverRejected, match="closed"): - driver.register(writer.payload.address, SLOT, RegionLayout(SLOT, 1)) + driver.register(writer.memory.address, SLOT, RegionLayout(SLOT, 1)) def test_memory_driver_is_a_store_driver(self) -> None: driver: StoreDriver = self.make() assert isinstance(driver, StoreDriver) + # The conformance harness itself. -# Drivers that break exactly one rule. Each proves that the check guarding the rule fails -# against it, so a weakened check is caught here before a real driver meets it. - - -class _AcceptsDuplicateKeys(MemoryDriver): - def _check_unique_keys(self, rows: Sequence[RowIO]) -> None: - pass - - -class _TouchesMemoryBeforeChecking(MemoryDriver): - def get(self, rows: Sequence[RowIO], hint: Hint = None) -> MemoryOp: - for io in rows: - ctypes.memset(io.payload.address, 0, io.payload.nbytes) - return super().get(rows, hint) - - -class _MissTouchesTheDestination(MemoryDriver): - def _get_row(self, io: RowIO) -> RowResult: - result = super()._get_row(io) - if result.status is RowStatus.MISS: - ctypes.memset(io.payload.address, 0, io.payload.nbytes) - return result - - -class _CallsAWrongSizeAMiss(MemoryDriver): - def _get_row(self, io: RowIO) -> RowResult: - result = super()._get_row(io) - return MISS if result.status is RowStatus.FAULT else result - - -class _OverwritesExistingObjects(MemoryDriver): - def _put_row(self, io: RowIO) -> RowResult: - payload = ctypes.string_at(io.payload.address, io.payload.nbytes) - self.store.overwrite(io.key, payload, self._trailer_for(io, payload)) - return super()._put_row(io) - - -class _IgnoresTheRowLimit(MemoryDriver): - def _check_row_limit(self, rows: Sequence[RowIO]) -> None: - if not rows: - raise DriverRejected("an operation needs at least one row") - - -class _IgnoresTheDirection(MemoryDriver): - def _check_direction(self, kind: Kind) -> None: - pass - - -class _NeverEndsAccess(MemoryDriver): - def _execute(self, op: MemoryOp) -> None: - super()._execute(op) - with self._lock: - op.access_ended = False # the evidence is lost again + def _harness(self, driver: MemoryDriver) -> Harness: + return Harness(TestMemoryDriverConformance(), driver, TIMEOUT) + def test_harness_keeps_a_buffer_whose_registration_escaped(self) -> None: + driver = self.make() + h = self._harness(driver) + driver.fail_next_register("timed out after the request", escaped=True) + with pytest.raises(RegistrationError): + h.endpoint(["A"]) + assert len(h.endpoints) == 1 and h.endpoints[0].handle is not None + assert h.cleanup() == [] + assert driver.live_registrations() == [] -class _ChangesItsCapabilities(MemoryDriver): - _flip = False + def test_harness_retains_an_endpoint_it_could_not_unregister_and_retries(self) -> None: + driver = self.make() + h = self._harness(driver) + made = h.endpoint(["A"]) + assert made.handle is not None + driver.fail_next_unregister(made.handle, "busy") + problems = h.cleanup() + assert problems and not h.done + assert made in UNRELEASED_ENDPOINTS + assert driver.live_registrations() == [made.handle] + assert h.cleanup() == [] and h.done + assert driver.live_registrations() == [] - @property - def capabilities(self) -> DriverCapabilities: - self._flip = not self._flip - return DriverCapabilities( - get=True, put=True, cancel=self._flip, probe=True, max_rows_per_op=None + def test_harness_retains_endpoints_when_drain_does_not_finish(self) -> None: + driver = type("DrainSaysNo", (self.driver_cls,), {"drain": lambda self, budget_s: False})( + MemoryStore() ) + h = self._harness(driver) + made = h.endpoint(["A"]) + problems = h.cleanup() + assert problems == ["drain did not return True"] and not h.done + assert made in UNRELEASED_ENDPOINTS + + def test_retained_pool_memory_outlives_its_suite(self) -> None: + """A pool endpoint retained for an operation that never proved its end keeps the pool + itself alive, not just an address into it.""" + + class ControlledPool(TestPoolDriverConformance): + def pool_driver(self) -> MemoryDriver: + return MemoryDriver(MemoryStore(), controlled=True) + + suite = ControlledPool() + driver = suite.make_driver() + assert isinstance(driver, MemoryDriver) + pool_alive = weakref.ref(suite.pool) + h = Harness(suite, driver, TIMEOUT) + made = h.endpoint(["A"]) + op = driver.get(made.rows()) + h.ops.append(op) + driver.transfer(op) + driver.publish_results(op) + driver.mark_unprovable(op) + assert h.cleanup() == [] + assert made in UNRELEASED_ENDPOINTS + del suite, h, made + gc.collect() + assert pool_alive() is not None, "the pool was freed while the driver may still touch it" + assert driver.live_registrations() + + def test_cleanup_retry_releases_pool_memory_once(self) -> None: + class DrainFailsOnce(MemoryDriver): + drains = 0 + + def drain(self, budget_s: float) -> bool: + self.drains += 1 + return super().drain(budget_s) and self.drains > 1 + + class CountingPool(TestPoolDriverConformance): + releases: list[object] + + def pool_driver(self) -> MemoryDriver: + return DrainFailsOnce(MemoryStore()) + + def release(self, memory: object) -> None: + self.releases.append(memory) + + suite = CountingPool() + suite.releases = [] + driver = suite.make_driver() + h = Harness(suite, driver, TIMEOUT) + made = h.endpoint(["A"]) + assert h.cleanup() == ["drain did not return True"] + assert suite.releases == [made.memory] + assert made in UNRELEASED_ENDPOINTS + assert h.cleanup() == [] and h.done + assert suite.releases == [made.memory] # once, not twice - -class _ProbeTakesItsTime(MemoryDriver): - def probe( - self, keys: Sequence[str], budget_s: float, hint: Hint = None - ) -> tuple[ProbeStatus, ...]: - time.sleep(budget_s + 0.8) - return super().probe(keys, budget_s, hint) - - -class _KeepsAcceptingWhileDraining(MemoryDriver): - def _check_not_draining(self) -> None: - pass - - -class _ClosesAnyTime(MemoryDriver): - def _check_closable(self) -> None: - pass + def test_harness_retains_endpoints_of_an_unprovable_operation_without_complaint(self) -> None: + driver = self.make(controlled=True) + h = self._harness(driver) + reader = h.endpoint(["A"]) + op = driver.get(reader.rows()) + h.ops.append(op) + driver.transfer(op) + driver.publish_results(op) + driver.mark_unprovable(op) + assert h.cleanup() == [] + assert reader in UNRELEASED_ENDPOINTS and not h.done + assert reader.handle in driver.live_registrations() -class _UnregisterKeepsTheRange(MemoryDriver): - def unregister(self, handle: RegionHandle) -> None: - pass +# Drivers that break one rule only the in-memory tests can observe. class _UnregistersBusyRegions(MemoryDriver): @@ -902,25 +931,11 @@ def _check_region_idle(self, registration: object) -> None: class _CallsATimeoutAMiss(MemoryDriver): - def _move(self, kind: Kind, io: RowIO) -> RowResult: - try: - return super()._move(kind, io) - finally: - pass - def _get_row(self, io: RowIO) -> RowResult: try: return super()._get_row(io) except StoreError: - return MISS - - -class _LetsRaisesEscape(MemoryDriver): - def _move(self, kind: Kind, io: RowIO) -> RowResult: - try: - return self._put_row(io) if kind == "put" else self._get_row(io) - except StoreError as error: - return faulted(str(error)) + return RowResult(RowStatus.MISS) class _SkipsTheIntegrityCheck(MemoryDriver): @@ -928,32 +943,9 @@ def _verify(self, io: RowIO, stored: StoredObject) -> str | None: return None -class _CallsNotReadyAFault(MemoryDriver): - def _get_row(self, io: RowIO) -> RowResult: - result = super()._get_row(io) - return faulted("not ready") if result.status is RowStatus.NOT_READY else result - - -class _CallsDeclinedAFault(MemoryDriver): - def _put_row(self, io: RowIO) -> RowResult: - result = super()._put_row(io) - return faulted("no space") if result.status is RowStatus.DECLINED else result - - -class _ProbeIgnoresTheBudget(MemoryDriver): - def probe( - self, keys: Sequence[str], budget_s: float, hint: Hint = None - ) -> tuple[ProbeStatus, ...]: - return tuple(self.store.presence(key) for key in keys) - - class _ProbesWithoutLimit(MemoryDriver): - def probe( - self, keys: Sequence[str], budget_s: float, hint: Hint = None - ) -> tuple[ProbeStatus, ...]: - with self._lock: - self._probes = self._probes[: MAX_PROBES_IN_FLIGHT - 1] - return super().probe(keys, budget_s, hint) + def _probe_slot_free(self) -> bool: + return True # the workers stay on the books, so they are still joined class _ReportsAccessEndedEarly(MemoryDriver): @@ -966,6 +958,16 @@ def _check_hint(self, hint: Hint) -> None: pass +class _AcceptsOverlappingSpans(MemoryDriver): + def _check_disjoint(self, kind: Kind, rows: Sequence[RowIO]) -> None: + pass + + +class _TakesAnyRowShape(MemoryDriver): + def _check_span_count(self, rows: Sequence[RowIO]) -> None: + pass + + class _ExecutesWithTheControls(MemoryDriver): def _execute(self, op: MemoryOp) -> None: self.transfer(op) @@ -978,85 +980,8 @@ def _drain_can_return(self) -> bool: return self._quiet() -LIMITED = DriverCapabilities(get=True, put=True, cancel=True, probe=True, max_rows_per_op=1) -GET_ONLY = DriverCapabilities(get=True, put=False, cancel=True, probe=True, max_rows_per_op=None) - -# rule -> (driver breaking it, the capabilities to run with, the conformance check that must fail) -BROKEN_CONFORMANCE = { - "a key appears once per operation": ( - _AcceptsDuplicateKeys, - DEFAULT_CAPABILITIES, - "test_a_key_appears_at_most_once_per_operation", - ), - "rejection happens before any effect": ( - _TouchesMemoryBeforeChecking, - DEFAULT_CAPABILITIES, - "test_rows_outside_a_live_registration_are_rejected_before_any_effect", - ), - "a miss leaves the destination alone": ( - _MissTouchesTheDestination, - DEFAULT_CAPABILITIES, - "test_get_reports_hit_hit_miss_and_leaves_the_miss_untouched", - ), - "a wrong size is a fault": ( - _CallsAWrongSizeAMiss, - DEFAULT_CAPABILITIES, - "test_an_object_of_the_wrong_size_is_a_fault_not_a_miss", - ), - "a put never damages an existing object": ( - _OverwritesExistingObjects, - DEFAULT_CAPABILITIES, - "test_a_put_under_an_existing_key_leaves_the_object_whole", - ), - "the row limit is enforced": ( - _IgnoresTheRowLimit, - LIMITED, - "test_rows_beyond_the_declared_limit_are_rejected", - ), - "an undeclared direction is rejected": ( - _IgnoresTheDirection, - GET_ONLY, - "test_an_undeclared_direction_is_rejected", - ), - "access ends with evidence": ( - _NeverEndsAccess, - DEFAULT_CAPABILITIES, - "test_results_are_stable_and_access_ends_with_evidence", - ), - "capabilities are fixed": ( - _ChangesItsCapabilities, - DEFAULT_CAPABILITIES, - "test_capabilities_do_not_change", - ), - "probe returns within its budget": ( - _ProbeTakesItsTime, - DEFAULT_CAPABILITIES, - "test_probe_returns_within_its_budget", - ), - "probe answers unknown for a blocked key": ( - _ProbeIgnoresTheBudget, - DEFAULT_CAPABILITIES, - "test_probe_answers_unknown_for_a_key_the_store_does_not_answer_in_time", - ), - "drain is one-way": ( - _KeepsAcceptingWhileDraining, - DEFAULT_CAPABILITIES, - "test_drain_is_one_way", - ), - "close needs a true drain": ( - _ClosesAnyTime, - DEFAULT_CAPABILITIES, - "test_close_needs_a_true_drain_and_no_live_registrations", - ), - "unregister frees the range": ( - _UnregisterKeepsTheRange, - DEFAULT_CAPABILITIES, - "test_unregister_is_idempotent_and_frees_the_range", - ), -} - # rule -> (driver breaking it, the memory-driver check that must fail) -BROKEN_MEMORY = { +BROKEN_MEMORY: dict[str, tuple[type[MemoryDriver], str]] = { "unregister needs an idle region": ( _UnregistersBusyRegions, "test_unregister_is_refused_while_an_operation_may_access_the_region", @@ -1065,38 +990,34 @@ def _drain_can_return(self) -> bool: _CallsATimeoutAMiss, "test_a_store_timeout_is_a_fault_not_a_miss", ), - "a recorded operation never raises": ( - _LetsRaisesEscape, - "test_a_raise_after_the_operation_is_recorded_becomes_a_row_fault", - ), "integrity is verified on every get": ( _SkipsTheIntegrityCheck, "test_a_corrupted_object_is_a_fault_and_never_lands", ), - "an unfinished write is not ready": ( - _CallsNotReadyAFault, - "test_a_key_held_by_an_unfinished_write_is_not_ready_for_get_put_and_probe", + "probes in flight are capped": ( + _ProbesWithoutLimit, + "test_probes_in_flight_are_capped", ), - "no room is declined": ( - _CallsDeclinedAFault, - "test_a_row_the_store_has_no_room_for_is_declined_not_faulted", + "access does not end early": ( + _ReportsAccessEndedEarly, + "test_queries_do_not_wait_on_an_operation_inside_the_store", ), "an unreadable hint is refused": ( _AcceptsHintsItCannotRead, "test_memory_driver_rejects_a_hint_it_cannot_read_before_any_effect", ), + "overlapping spans are refused": ( + _AcceptsOverlappingSpans, + "test_memory_driver_rejects_overlapping_spans_within_one_get", + ), + "an undeclared multi-span row is refused": ( + _TakesAnyRowShape, + "test_a_single_span_driver_rejects_a_multi_span_row", + ), "a submission racing a drain keeps its handle": ( _ExecutesWithTheControls, "test_a_drain_racing_a_submission_leaves_a_cancelled_operation_with_its_handle", ), - "probes in flight are capped": ( - _ProbesWithoutLimit, - "test_probes_in_flight_are_capped", - ), - "access does not end early": ( - _ReportsAccessEndedEarly, - "test_queries_do_not_wait_on_an_operation_inside_the_store", - ), "drain returns at once on an unprovable operation": ( _DrainWaitsForEveryone, "test_drain_returns_false_at_once_even_while_another_operation_is_still_provable", @@ -1106,6 +1027,18 @@ def _drain_can_return(self) -> bool: CHECK_FAILED = (AssertionError, pytest.fail.Exception) +class UnexpectedSkip(Exception): + """A paired check skipped; not an expected failure, so no ``pytest.raises`` swallows it.""" + + +def _run_check(suite: object, check: str) -> None: + """Runs one check; a skip is a failure here, because a pairing that skips proves nothing.""" + try: + getattr(suite, check)() + except pytest.skip.Exception as skipped: + raise UnexpectedSkip(f"{check} skipped instead of running: {skipped}") from None + + def _conformance(driver_cls: type[MemoryDriver], caps: DriverCapabilities) -> DriverConformance: suite = type( "Suite", @@ -1120,9 +1053,9 @@ def test_each_conformance_check_passes_the_driver_and_fails_the_one_breaking_its rule: str, ) -> None: broken, capabilities, check = BROKEN_CONFORMANCE[rule] - getattr(_conformance(MemoryDriver, capabilities), check)() + _run_check(_conformance(MemoryDriver, capabilities), check) with pytest.raises(CHECK_FAILED): - getattr(_conformance(broken, capabilities), check)() + _run_check(_conformance(broken, capabilities), check) @pytest.mark.parametrize("rule", sorted(BROKEN_MEMORY)) @@ -1130,7 +1063,17 @@ def test_each_memory_driver_check_fails_the_driver_breaking_its_rule(rule: str) broken, check = BROKEN_MEMORY[rule] suite = type("Suite", (TestMemoryDriver,), {"driver_cls": broken})() with pytest.raises(CHECK_FAILED): - getattr(suite, check)() + _run_check(suite, check) + + +def test_a_skip_on_the_broken_side_is_not_an_expected_failure() -> None: + no_probe = DriverCapabilities( + get=True, put=True, cancel=True, probe=False, max_rows_per_op=None + ) + suite = _conformance(MemoryDriver, no_probe) + with pytest.raises(UnexpectedSkip): + with pytest.raises(CHECK_FAILED): # the shape the pairing test uses; must not catch it + _run_check(suite, "test_probe_reports_hit_and_miss") def test_each_broken_driver_breaks_one_rule_and_names_an_existing_check() -> None: diff --git a/tests/unittest/_torch/disaggregation/blob/test_testing_isolation.py b/tests/unittest/_torch/disaggregation/blob/test_testing_isolation.py new file mode 100644 index 000000000000..92eb578ea560 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_testing_isolation.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The test tooling ships in the wheel for driver teams; no engine module may import it.""" + +from __future__ import annotations + +import ast +import pathlib + +import pytest + +import tensorrt_llm._torch.disaggregation.testing as testing + +pytestmark = pytest.mark.cpu_only + +TESTING = "tensorrt_llm._torch.disaggregation.testing" +TESTING_DIR = pathlib.Path(testing.__file__).resolve().parent +PACKAGE_DIR = TESTING_DIR.parents[2] # .../tensorrt_llm + + +def _package_of(path: pathlib.Path) -> str: + """The package a module at ``path`` resolves relative imports against.""" + relative = path.resolve().relative_to(PACKAGE_DIR.parent).with_suffix("") + parts = list(relative.parts) + if parts[-1] != "__init__": + parts.pop() + else: + parts.pop() + return ".".join(parts) if path.name == "__init__.py" else ".".join(parts) + + +def _testing_imports(source: str, package: str) -> list[str]: + """Absolute names of everything ``source`` imports from the testing package, deferred and + relative imports included.""" + found: list[str] = [] + for node in ast.walk(ast.parse(source)): + if isinstance(node, ast.Import): + found.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + module = node.module or "" + if node.level: + anchor = package.rsplit(".", node.level - 1)[0] if node.level > 1 else package + module = f"{anchor}.{module}" if module else anchor + found.append(module) + found.extend(f"{module}.{alias.name}" for alias in node.names) + return [name for name in found if name == TESTING or name.startswith(TESTING + ".")] + + +def test_no_engine_module_imports_the_testing_package() -> None: + offenders = {} + for path in PACKAGE_DIR.rglob("*.py"): + if TESTING_DIR in path.resolve().parents: + continue + source = path.read_text(encoding="utf-8", errors="replace") + if "testing" not in source: + continue # cannot name the package without the word + found = _testing_imports(source, _package_of(path)) + if found: + offenders[str(path.relative_to(PACKAGE_DIR))] = found + assert offenders == {} + + +def test_the_import_scan_sees_every_kind_of_import() -> None: + source = ( + "from typing import TYPE_CHECKING\n" + "if TYPE_CHECKING:\n" + " from ..testing import driver_conformance\n" + "def later():\n" + " import tensorrt_llm._torch.disaggregation.testing.memory_store\n" + " from ..testing.host_memory import HostBuffer\n" + " from . import keys\n" + ) + found = _testing_imports(source, "tensorrt_llm._torch.disaggregation.blob") + assert f"{TESTING}.driver_conformance" in found + assert f"{TESTING}.memory_store" in found + assert f"{TESTING}.host_memory.HostBuffer" in found + assert len(found) == 5 # the three modules, two members; nothing from the blob import + assert _package_of(PACKAGE_DIR / "_torch" / "disaggregation" / "blob" / "store.py") == ( + "tensorrt_llm._torch.disaggregation.blob" + ) + assert _package_of(PACKAGE_DIR / "_torch" / "disaggregation" / "__init__.py") == ( + "tensorrt_llm._torch.disaggregation" + ) From e15391685159ad5f45cc4eac13b85e203814cfff Mon Sep 17 00:00:00 2001 From: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:42:40 +0800 Subject: [PATCH 5/6] [None][test] End the blocked-probe test's wait at the gate instead of on the clock The in-memory driver's probe wait is a small hook; the test overrides it to return once the worker is stuck at the gate, so the expected answers no longer depend on thread scheduling. Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> --- .../_torch/disaggregation/testing/memory_store.py | 7 ++++++- .../_torch/disaggregation/blob/test_memory_driver.py | 11 ++++++++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/testing/memory_store.py b/tensorrt_llm/_torch/disaggregation/testing/memory_store.py index 5b1daac592eb..746be8114816 100644 --- a/tensorrt_llm/_torch/disaggregation/testing/memory_store.py +++ b/tensorrt_llm/_torch/disaggregation/testing/memory_store.py @@ -347,7 +347,7 @@ def ask() -> None: self._probes.append(_Probe(worker, finished)) self._probe_workers.append(worker) worker.start() - if finished.wait(budget_s): + if self._wait_probe(finished, budget_s): worker.join() # it is exiting; leave no thread behind a call that got its answer if errors: raise errors[0] @@ -552,6 +552,11 @@ def _check_closable(self) -> None: def _probe_slot_free(self) -> bool: return len(self._probes) < MAX_PROBES_IN_FLIGHT + def _wait_probe(self, finished: threading.Event, budget_s: float) -> bool: + """Waits for the probe worker up to the budget; True when it finished in time. A test + that must not depend on scheduling overrides this to end the wait on its own signal.""" + return finished.wait(budget_s) + def _start_draining(self) -> None: self._draining = True diff --git a/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py index 448e822039cb..c8ed766c4496 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py +++ b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py @@ -557,7 +557,16 @@ def test_an_ineffective_cancel_changes_nothing(self) -> None: def test_a_blocked_probe_holds_drain_until_the_store_answers(self) -> None: store = GatedStore(held={KEY_A}) - driver = self.make(store) + base = self.driver_cls + + class BudgetEndsAtTheGate(base): # type: ignore[misc,valid-type] + # The wait ends once the worker is stuck at the gate, not after a wall-clock budget: + # by then B is answered and A, C are not, whatever the scheduler did. + def _wait_probe(self, finished: threading.Event, budget_s: float) -> bool: + assert store.waiting.wait(TIMEOUT), "the probe worker never reached the gate" + return False # as if the budget had run out + + driver = BudgetEndsAtTheGate(store) self._publish(driver, {"B": b"bbbbbbbb"}) store.gate.clear() try: From 09fa576dcece2a8685447dfdea5688eed7d41166 Mon Sep 17 00:00:00 2001 From: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> Date: Mon, 19 Oct 2026 19:57:34 +0800 Subject: [PATCH 6/6] [None][feat] Carry the memory kind through the driver boundary and make more driver checks bite A DEVICE region is rejected unless the driver declares device_memory, a MISS leaves its destination untouched, and cancel, put sources, absent keys and row limits now have checks a broken driver fails. The in-memory driver counts its background workers in drain and concludes an operation whose thread cannot start. Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> --- .../_torch/disaggregation/base/memory.py | 28 ++ .../_torch/disaggregation/blob/__init__.py | 4 +- .../_torch/disaggregation/blob/store.py | 41 ++- .../_torch/disaggregation/testing/__init__.py | 9 + .../disaggregation/testing/broken_drivers.py | 160 ++++++++- .../testing/driver_conformance.py | 308 ++++++++++++++++-- .../disaggregation/testing/memory_store.py | 65 +++- .../disaggregation/blob/test_blob_format.py | 3 + .../disaggregation/blob/test_memory_driver.py | 134 +++++--- .../blob/test_testing_isolation.py | 7 +- 10 files changed, 649 insertions(+), 110 deletions(-) create mode 100644 tensorrt_llm/_torch/disaggregation/base/memory.py diff --git a/tensorrt_llm/_torch/disaggregation/base/memory.py b/tensorrt_llm/_torch/disaggregation/base/memory.py new file mode 100644 index 000000000000..34cd71631661 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/base/memory.py @@ -0,0 +1,28 @@ +# 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 kinds of memory a cache backend can access.""" + +import enum + +__all__ = ["MemoryKind"] + + +class MemoryKind(enum.Enum): + """Memory a backend can access: HOST is CPU-addressable; DEVICE is the GPU holding the cache. + + Each region of memory is of one kind.""" + + HOST = "host" + DEVICE = "device" diff --git a/tensorrt_llm/_torch/disaggregation/blob/__init__.py b/tensorrt_llm/_torch/disaggregation/blob/__init__.py index 5e7ae28d4d83..efab37bb6b99 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/blob/__init__.py @@ -3,6 +3,6 @@ """The blob backend: KV rows shared through a byte store. So far this package holds the private driver boundary (``store.py``) and the key encoding -(``keys.py``). Nothing here imports the public backend contract or an optional store SDK at -import time. +(``keys.py``). Nothing here imports an optional store SDK at import time, and the only thing +shared with the public backend contract is ``MemoryKind``. """ diff --git a/tensorrt_llm/_torch/disaggregation/blob/store.py b/tensorrt_llm/_torch/disaggregation/blob/store.py index 90385b37496f..acf2d60fede9 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/store.py +++ b/tensorrt_llm/_torch/disaggregation/blob/store.py @@ -5,8 +5,8 @@ A driver moves whole rows between registered memory and one store. It answers two questions separately: what each row's logical result is (:meth:`StoreDriver.poll`), and whether an operation can still touch the caller's memory (:meth:`StoreDriver.access_ended`). Neither answer implies the -other. Nothing here belongs to the public backend contract in ``base/``; the backend translates -between the two. +other. Nothing here belongs to the public backend contract in ``base/``, from which only +``MemoryKind`` is shared; the backend translates between the two. """ from __future__ import annotations @@ -16,6 +16,8 @@ from dataclasses import dataclass from typing import Any, Protocol, runtime_checkable +from ..base.memory import MemoryKind + Hint = Any """The caller's hint for where a get or probe should look, handed to the driver unchanged. A hint is an immutable snapshot: the driver may keep a reference to it for as long as it needs and never @@ -69,8 +71,10 @@ class RowStatus(enum.Enum): store's "already exists" counts only once that object is committed.""" MISS = "miss" - """For a get: the store holds no object under the key. Only a driver that declares - ``absent_is_miss`` can tell this apart from a failure; another reports FAULT for such a key.""" + """For a get: the store reported that it holds no object under the key, and the row's memory + was not touched. Every driver answers MISS only for a reported absence and FAULT for any + failure; a driver that declares ``absent_is_miss`` adds that an absent key always concludes + MISS.""" NOT_READY = "not_ready" """The key is held by a write that has not finished: a get delivered nothing, a put published @@ -146,11 +150,14 @@ class RegionLayout: """How a registered region is divided: ``slot_count`` slots of ``slot_bytes``, back to back. ``label`` is an opaque layout label, equal for regions laid out alike; it is not an identity. + ``memory_kind`` says whether the region is host or device memory; a transport registers the + two differently. """ slot_bytes: int slot_count: int label: str = "" + memory_kind: MemoryKind = MemoryKind.HOST def __post_init__(self) -> None: if self.slot_bytes <= 0 or self.slot_count <= 0: @@ -228,11 +235,10 @@ class DriverCapabilities: max_rows_per_op: Rows one ``put``/``get`` accepts, or ``None`` for no limit. hints: ``get`` and ``probe`` read a caller's hint; without it the backend rejects a submission carrying a hint before dispatch and the driver is called without one. - absent_is_miss: A get of a key the store does not hold concludes MISS. Without it the - store cannot tell absence from failure, such a get may conclude FAULT, and the - backend's fault accounting counts misses too. - device_memory: The driver can register regions of device memory; without it every region - handed to it is host memory. + absent_is_miss: A get of a key the store does not hold always concludes MISS. Without + it the store may not tell absence from failure, and such a get may conclude FAULT. + device_memory: The driver registers regions whose ``memory_kind`` is DEVICE; without it + every region handed to it is host memory and it rejects a DEVICE one. multi_span: ``put`` and ``get`` take rows of more than one span; without it every row handed to it has exactly one. """ @@ -293,11 +299,14 @@ def capabilities(self) -> DriverCapabilities: def register(self, address: int, nbytes: int, layout: RegionLayout) -> RegionHandle: """Register one region the driver may read and write. - The region is host memory unless the driver declares ``device_memory``, in which case the - caller may also hand it device regions. + ``layout.memory_kind`` says whether the region is host or device memory; a DEVICE region + is handed only to a driver that declares ``device_memory``. The kind is checked at this + entry: a driver without ``device_memory`` rejects a DEVICE region before the store is + asked or the memory touched, so the mismatch never reaches a registration or a transfer. Raises: - DriverRejected: Bad geometry or overlap with a live registration; nothing escaped. + DriverRejected: Bad geometry, overlap with a live registration, or a DEVICE region + for a driver without ``device_memory``; nothing escaped. RegistrationError: The store refused; ``handle`` is set when the store may have registered the region anyway. """ @@ -316,13 +325,13 @@ def put(self, rows: Sequence[RowIO]) -> OpHandle: """Publish ``rows``; may return before any byte moved. The caller guarantees that the rows' spans hold the payload and stay unchanged until - ``access_ended`` is true. + ``access_ended`` is true, and hands a row of more than one span only to a driver that + declares ``multi_span``; a driver may treat an unexpected one as a programming error. Raises: DriverRejected: ``put`` is not declared, the driver is draining or closed, a key - appears twice in ``rows``, ``rows`` exceed ``max_rows_per_op``, a row has more - than one span without ``multi_span``, or a span lies outside a live registration; - nothing escaped. + appears twice in ``rows``, ``rows`` exceed ``max_rows_per_op``, or a span lies + outside a live registration; nothing escaped. """ ... diff --git a/tensorrt_llm/_torch/disaggregation/testing/__init__.py b/tensorrt_llm/_torch/disaggregation/testing/__init__.py index 511decccb36a..6ad7306bcf28 100644 --- a/tensorrt_llm/_torch/disaggregation/testing/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/testing/__init__.py @@ -11,3 +11,12 @@ test guards that. It provides no public configuration: the in-memory driver cannot be selected in a deployment. """ + +import pytest + +# The checks are library modules, which pytest does not rewrite on its own; registered here, a +# failing check shows the values it compared rather than a bare AssertionError. +pytest.register_assert_rewrite( + "tensorrt_llm._torch.disaggregation.testing.driver_conformance", + "tensorrt_llm._torch.disaggregation.testing.broken_drivers", +) diff --git a/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py b/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py index cbe1a25550bc..69235c6e2b1d 100644 --- a/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py +++ b/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py @@ -12,24 +12,29 @@ import ctypes import time from collections.abc import Sequence +from typing import Any from tensorrt_llm._torch.disaggregation.blob.store import ( DriverCapabilities, DriverRejected, Hint, + OpHandle, ProbeStatus, RegionHandle, + RegionLayout, RowIO, RowResult, RowStatus, + Span, StoreError, faulted, ) -from .memory_store import DEFAULT_CAPABILITIES, Kind, MemoryDriver, MemoryOp, _scatter +from .memory_store import DEFAULT_CAPABILITIES, Kind, MemoryDriver, MemoryOp, MemoryStore, _scatter _MISS = RowResult(RowStatus.MISS) _READY = RowResult(RowStatus.READY) +_CANCELLED = RowResult(RowStatus.CANCELLED) class AcceptsDuplicateKeys(MemoryDriver): @@ -45,6 +50,13 @@ def get(self, rows: Sequence[RowIO], hint: Hint = None) -> MemoryOp: return super().get(rows, hint) +class ChecksOnlyTheStartOfASpan(MemoryDriver): + def _registered(self, span: Span) -> bool: + return any( + r.live and r.span.address <= span.address < r.span.end for r in self._registrations + ) + + class MissTouchesTheDestination(MemoryDriver): def _get_row(self, io: RowIO) -> RowResult: result = super()._get_row(io) @@ -62,12 +74,36 @@ def _get_row(self, io: RowIO) -> RowResult: return super()._get_row(io) +class TruncatesIntoMultiSpanRows(MemoryDriver): + def _get_row(self, io: RowIO) -> RowResult: + stored = self.store.get(io.key) + if len(io.spans) > 1 and stored is not None and len(stored.payload) > io.nbytes: + _scatter(io.spans, stored.payload[: io.nbytes]) + return _READY + return super()._get_row(io) + + class CallsAWrongSizeAMiss(MemoryDriver): def _get_row(self, io: RowIO) -> RowResult: result = super()._get_row(io) return _MISS if result.status is RowStatus.FAULT else result +class CallsATimeoutAMiss(MemoryDriver): + def _get_row(self, io: RowIO) -> RowResult: + try: + return super()._get_row(io) + except StoreError: + return _MISS + + +class ClobbersItsSource(MemoryDriver): + def _put_row(self, io: RowIO) -> RowResult: + result = super()._put_row(io) + _scatter(io.spans, bytes(io.nbytes)) + return result + + class OverwritesExistingObjects(MemoryDriver): def _put_row(self, io: RowIO) -> RowResult: payload = b"".join(ctypes.string_at(span.address, span.nbytes) for span in io.spans) @@ -103,6 +139,11 @@ def _check_direction(self, kind: Kind) -> None: pass +class RegistersAnyMemoryKind(MemoryDriver): + def _check_region_kind(self, layout: RegionLayout) -> None: + pass + + class NeverEndsAccess(MemoryDriver): def _execute(self, op: MemoryOp) -> None: super()._execute(op) @@ -123,6 +164,7 @@ def capabilities(self) -> DriverCapabilities: probe=True, max_rows_per_op=None, absent_is_miss=True, + device_memory=True, multi_span=True, ) @@ -158,6 +200,11 @@ def _execute(self, op: MemoryOp) -> None: super()._execute(op) +class UnregistersBusyRegions(MemoryDriver): + def _check_region_idle(self, registration: object) -> None: + pass + + class CallsNotReadyAFault(MemoryDriver): def _get_row(self, io: RowIO) -> RowResult: result = super()._get_row(io) @@ -178,6 +225,11 @@ def _move(self, kind: Kind, io: RowIO) -> RowResult: return faulted(str(error)) +class DrainWaitsForEveryone(MemoryDriver): + def _drain_can_return(self) -> bool: + return self._quiet() + + class KeepsAcceptingWhileDraining(MemoryDriver): def _check_not_draining(self) -> None: pass @@ -193,12 +245,59 @@ def unregister(self, handle: RegionHandle) -> None: pass +class CancelWaitsForTheTransfer(MemoryDriver): + def cancel(self, op: OpHandle) -> None: + with self._lock: + self._lock.wait_for(lambda: self._op(op).access_ended) + super().cancel(op) + + +class CancelsStartedRows(MemoryDriver): + def _cancel_unstarted(self, op: MemoryOp) -> None: + for row in op.rows: + if not row.published.concluded: + row.published = _CANCELLED + self._settle_untouched(op) + + +class CancelEndsAccessWithTheRowInTheStore(MemoryDriver): + """On cancel, a row the store holds is called FAULT and the access over, as if the transfer + had been torn down; the transfer in the store goes on regardless.""" + + def __init__(self, store: MemoryStore, **options: Any) -> None: + super().__init__(store, **options) + self._torn_down: set[str] = set() + + def cancel(self, op: OpHandle) -> None: + with self._lock: + recorded = self._op(op) + for row in recorded.rows: + if row.started and not row.published.concluded: + row.published = faulted("torn down") + self._torn_down.add(row.io.key) + recorded.access_ended = True + self._lock.notify_all() + + def _move(self, kind: Kind, io: RowIO) -> RowResult: + result = super()._move(kind, io) + return faulted("torn down") if io.key in self._torn_down else result + + LIMITED = DriverCapabilities( get=True, put=True, cancel=True, probe=True, max_rows_per_op=1, absent_is_miss=True ) GET_ONLY = DriverCapabilities( get=True, put=False, cancel=True, probe=True, max_rows_per_op=None, absent_is_miss=True ) +HOST_ONLY = DriverCapabilities( + get=True, + put=True, + cancel=True, + probe=True, + max_rows_per_op=None, + absent_is_miss=True, + multi_span=True, +) # rule -> (driver breaking it, the capabilities to run with, the conformance check that must fail) BROKEN_CONFORMANCE: dict[str, tuple[type[MemoryDriver], DriverCapabilities, str]] = { @@ -210,7 +309,12 @@ def unregister(self, handle: RegionHandle) -> None: "rejection happens before any effect": ( TouchesMemoryBeforeChecking, DEFAULT_CAPABILITIES, - "test_rows_outside_a_live_registration_are_rejected_before_any_effect", + "test_a_row_outside_every_registration_is_rejected_before_any_effect", + ), + "a span is checked to its end": ( + ChecksOnlyTheStartOfASpan, + DEFAULT_CAPABILITIES, + "test_a_span_leaving_or_outliving_its_registration_is_rejected", ), "a miss leaves the destination alone": ( MissTouchesTheDestination, @@ -222,11 +326,26 @@ def unregister(self, handle: RegionHandle) -> None: DEFAULT_CAPABILITIES, "test_an_object_larger_than_its_row_is_a_fault_not_a_truncated_ready", ), + "a larger object is not truncated into spans": ( + TruncatesIntoMultiSpanRows, + DEFAULT_CAPABILITIES, + "test_a_larger_object_is_not_truncated_into_a_multi_span_row", + ), "a smaller object is a fault": ( CallsAWrongSizeAMiss, DEFAULT_CAPABILITIES, "test_an_object_smaller_than_its_row_is_a_fault", ), + "a failed read is a fault": ( + CallsATimeoutAMiss, + DEFAULT_CAPABILITIES, + "test_a_failed_store_read_is_a_fault_not_a_miss", + ), + "a put leaves its source unchanged": ( + ClobbersItsSource, + DEFAULT_CAPABILITIES, + "test_put_then_get_moves_whole_rows_and_never_aliases_the_source", + ), "a put never damages an existing object": ( OverwritesExistingObjects, DEFAULT_CAPABILITIES, @@ -252,6 +371,11 @@ def unregister(self, handle: RegionHandle) -> None: GET_ONLY, "test_an_undeclared_direction_is_rejected", ), + "a device region needs device_memory": ( + RegistersAnyMemoryKind, + HOST_ONLY, + "test_a_device_region_is_rejected_unless_device_memory_is_declared", + ), "access ends with evidence": ( NeverEndsAccess, DEFAULT_CAPABILITIES, @@ -282,6 +406,11 @@ def unregister(self, handle: RegionHandle) -> None: DEFAULT_CAPABILITIES, "test_a_blocked_get_delays_neither_another_get_nor_the_non_blocking_calls", ), + "a busy region is not unregistered": ( + UnregistersBusyRegions, + DEFAULT_CAPABILITIES, + "test_a_blocked_get_delays_neither_another_get_nor_the_non_blocking_calls", + ), "an unfinished write is not ready": ( CallsNotReadyAFault, DEFAULT_CAPABILITIES, @@ -297,6 +426,11 @@ def unregister(self, handle: RegionHandle) -> None: DEFAULT_CAPABILITIES, "test_a_raise_after_the_operation_is_recorded_becomes_a_row_fault", ), + "drain returns at once on an unprovable operation": ( + DrainWaitsForEveryone, + DEFAULT_CAPABILITIES, + "test_an_unprovable_get_makes_drain_return_at_once_and_keeps_its_region", + ), "drain is one-way": ( KeepsAcceptingWhileDraining, DEFAULT_CAPABILITIES, @@ -310,6 +444,26 @@ def unregister(self, handle: RegionHandle) -> None: "unregister frees the range": ( UnregisterKeepsTheRange, DEFAULT_CAPABILITIES, - "test_unregister_is_idempotent_and_frees_the_range", + "test_unregister_frees_the_range_for_a_new_registration", + ), +} + +# The same, for checks that can only ask a driver that returns the handle before the store +# answers; the pairing test runs these against the asynchronous in-memory driver. +BROKEN_ASYNC_CONFORMANCE: dict[str, tuple[type[MemoryDriver], DriverCapabilities, str]] = { + "cancel never waits for the store": ( + CancelWaitsForTheTransfer, + DEFAULT_CAPABILITIES, + "test_cancel_of_a_get_inside_the_store_neither_waits_nor_lies", + ), + "a started row is never cancelled": ( + CancelsStartedRows, + DEFAULT_CAPABILITIES, + "test_cancel_of_a_get_inside_the_store_neither_waits_nor_lies", + ), + "access does not end with a row in the store": ( + CancelEndsAccessWithTheRowInTheStore, + DEFAULT_CAPABILITIES, + "test_cancel_of_a_get_inside_the_store_neither_waits_nor_lies", ), } diff --git a/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py b/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py index 9d4aacb92f92..a5d672296024 100644 --- a/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py +++ b/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py @@ -27,6 +27,7 @@ import pytest +from tensorrt_llm._torch.disaggregation.base.memory import MemoryKind from tensorrt_llm._torch.disaggregation.blob.keys import encode_key from tensorrt_llm._torch.disaggregation.blob.store import ( DriverRejected, @@ -47,6 +48,7 @@ from .host_memory import HostBuffer, Memory SLOT = 8 +ROW_LIMIT_CHECKED = 64 # a row limit this high is taken on trust, so the check stays small # Lender names are 54 bytes and start with a 16-byte scope digest shared by every name of a # deployment. A key built from a truncated name collides across these. @@ -93,12 +95,14 @@ def __init__( *, prefix: str, slot: int = SLOT, + memory_kind: MemoryKind = MemoryKind.HOST, ) -> None: self.driver = driver self.labels = list(labels) self.memory = memory self.prefix = prefix self.slot = slot + self.memory_kind = memory_kind self.handle = handle # A handle given up front belongs to whoever allocated the memory; this endpoint only # unregisters what it registered itself. @@ -108,7 +112,7 @@ def __init__( def register(self) -> None: """Registers the region. When the store may have kept it despite an error, the handle is recorded all the same, so the region can still be unregistered later.""" - layout = RegionLayout(self.slot, len(self.labels), "kv") + layout = RegionLayout(self.slot, len(self.labels), "kv", self.memory_kind) try: self.handle = self.driver.register(self.memory.address, self.memory.nbytes, layout) except RegistrationError as error: @@ -166,9 +170,24 @@ def __init__(self, suite: DriverConformance, driver: StoreDriver, timeout: float def key(self, label: str) -> str: return key_of(label, self.prefix) - def endpoint(self, labels: Sequence[str], *, fill: int = 0, slot: int = SLOT) -> Endpoint: - memory, handle = self.suite.allocate(slot * len(labels), fill) - endpoint = Endpoint(self.driver, labels, memory, handle, prefix=self.prefix, slot=slot) + def endpoint( + self, + labels: Sequence[str], + *, + fill: int = 0, + slot: int = SLOT, + memory_kind: MemoryKind = MemoryKind.HOST, + ) -> Endpoint: + memory, handle = self.suite.allocate(slot * len(labels), fill, memory_kind) + endpoint = Endpoint( + self.driver, + labels, + memory, + handle, + prefix=self.prefix, + slot=slot, + memory_kind=memory_kind, + ) self.endpoints.append(endpoint) # kept before the store is asked, whatever register does if handle is None: endpoint.register() @@ -186,13 +205,17 @@ def get(self, rows: Sequence[RowIO], hint: Hint = None) -> OpHandle: return op def publish(self, rows: dict[str, bytes], slot: int = SLOT) -> Endpoint: - """Puts ``rows`` (label to payload) and returns the writer endpoint.""" + """Puts ``rows`` (label to payload) and returns the writer endpoint, its source bytes + checked once access ended: a put reads its rows and never writes them.""" writer = self.endpoint(list(rows), slot=slot) for i, data in enumerate(rows.values()): writer.fill_row(i, data) op = self.put(writer.rows()) assert [r.status for r in self.settle(op)] == [RowStatus.READY] * len(rows) self.settle_access(op) + for i, data in enumerate(rows.values()): + assert writer.memory.read(i * slot, len(data)) == data, "a put changed its source" + assert writer.intact() return writer # Waiting and asserting. @@ -208,13 +231,15 @@ def settle_access(self, op: OpHandle) -> None: wait_until(lambda: self.driver.access_ended(op), "access to end", self.timeout) def assert_absent(self, result: RowResult, endpoint: Endpoint, index: int) -> None: - """What a get of a key nobody published must look like: MISS with the slot untouched - for a driver that can tell absence apart, otherwise anything but READY.""" - if self.driver.capabilities.absent_is_miss: - assert result.status is RowStatus.MISS, result + """What a get of a key nobody published must look like: MISS, or FAULT only for a driver + without ``absent_is_miss`` whose store could not tell absence from failure. Whichever + driver answers MISS, the slot is untouched.""" + allowed = [RowStatus.MISS] + if not self.driver.capabilities.absent_is_miss: + allowed.append(RowStatus.FAULT) + assert result.status in allowed, result + if result.status is RowStatus.MISS: assert endpoint.read_row(index) == bytes([0x2D]) * endpoint.slot - else: - assert result.status is not RowStatus.READY, result def needs( self, @@ -316,21 +341,30 @@ def _retain(self) -> None: class DriverConformance: """Contract checks for one driver. Subclass as ``Test...``, implement ``make_driver`` and - whichever scenario hooks the driver's store can stage.""" + whichever scenario hooks the driver's store can stage. + + The checks of a get prepare their objects through the driver's own put, so a driver that + declares no put skips them and passes the suite without one get being checked; its gets + count as covered only once its objects are staged some other way. + """ timeout = 5.0 registers_on_demand = True # endpoints register fresh memory; a pool driver sets False + handles_before_completion = False # put/get return their handle before the store answers def make_driver(self) -> StoreDriver: raise NotImplementedError # Memory and lifetime hooks. - def allocate(self, nbytes: int, fill: int) -> Allocation: - """Memory for one endpoint and, when the driver registered it up front, its handle. + def allocate(self, nbytes: int, fill: int, memory_kind: MemoryKind) -> Allocation: + """Memory of the given kind for one endpoint and, when the driver registered it up + front, its handle. The memory object keeps its backing allocation alive for as long as it exists: a check - that cannot prove access ended retains the endpoint, and nothing else, for good. + that cannot prove access ended retains the endpoint, and nothing else, for good. This + default stands in for either kind with host memory, which is what the in-memory driver + expects; a driver that really registers device memory allocates it here. """ return HostBuffer(nbytes, fill), None @@ -350,7 +384,9 @@ def make_blocked_probe(self) -> tuple[StoreDriver, str, Unblock] | None: def make_blocked_get(self) -> tuple[StoreDriver, str, threading.Event | None, Unblock] | None: """A driver whose store stalls every get of the returned key until the callable runs, - with an event set once a get is stalled, or ``None`` for no such signal.""" + with an event set once a get is stalled, or ``None`` for no such signal. A stalled get is + inside the store: its row has started, and its access cannot end before the callable + runs.""" return None def make_held_key(self) -> tuple[StoreDriver, str, Unblock] | None: @@ -370,6 +406,15 @@ def make_failing_presence(self) -> tuple[StoreDriver, str] | None: """A driver whose store fails the next presence query of the returned key.""" return None + def make_failing_get(self) -> tuple[StoreDriver, str] | None: + """A driver whose store fails every read of the returned key while a put of it works.""" + return None + + def make_unprovable_get(self) -> tuple[StoreDriver, Callable[[OpHandle], None]] | None: + """A driver and a callable that, given the handle of a get just submitted, concludes + it while the driver can never prove that its access ended.""" + return None + @contextlib.contextmanager def harness(self, driver: StoreDriver | None = None) -> Generator[Harness, None, None]: """A fresh driver (or ``driver``) with cleanup afterwards. A failing check keeps its own @@ -484,7 +529,12 @@ def test_a_multi_span_row_is_stored_and_fetched_back_to_back(self) -> None: assert got.status is RowStatus.READY assert whole.read_row(0) == b"aaaacccc" - parts = h.endpoint(["X", "Y", "Z"], slot=4, fill=0x2D) + kind = ( + MemoryKind.DEVICE + if h.driver.capabilities.device_memory and self.registers_on_demand + else MemoryKind.HOST + ) + parts = h.endpoint(["X", "Y", "Z"], slot=4, fill=0x2D, memory_kind=kind) back = RowIO(h.key("M"), (parts.span(2), parts.span(0))) # given order, not address (got,) = h.settle(h.get([back])) assert got.status is RowStatus.READY @@ -517,18 +567,58 @@ def test_a_row_may_cover_parts_of_slots_across_regions(self) -> None: assert split.read_row(0) == b"---WXY----xyzUV-" # the first 5 bytes land first assert left.intact() and right.intact() and whole.intact() and split.intact() + def test_a_larger_object_is_not_truncated_into_a_multi_span_row(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True, multi_span=True) + h.publish({"A": b"aaaaaaaa"}) + parts = h.endpoint(["X", "Y"], slot=3, fill=0x2D) + row = RowIO(h.key("A"), (parts.span(0), parts.span(1))) # 6 bytes for an 8-byte object + (got,) = h.settle(h.get([row])) + assert got.status is RowStatus.FAULT, got + assert parts.intact() + + def test_a_failed_store_read_is_a_fault_not_a_miss(self) -> None: + """A store whose reads fail must not look empty, or it would never be isolated.""" + staged = self.make_failing_get() + if staged is None: + pytest.skip("this driver cannot stage a failing read") + driver, key = staged + with self.harness(driver) as h: + h.needs(put=True, get=True) + writer = h.endpoint(["F"]) + writer.fill_row(0, b"ffffffff") + (put,) = h.settle(h.put([RowIO(key, (writer.span(0),))])) + assert put.status is RowStatus.READY + reader = h.endpoint(["F"], fill=0x2D) + (got,) = h.settle(h.get([RowIO(key, (reader.span(0),))])) + assert got.status is RowStatus.FAULT, got + assert reader.intact() + # Rejections leave no trace. - def test_rows_outside_a_live_registration_are_rejected_before_any_effect(self) -> None: + def test_a_row_outside_every_registration_is_rejected_before_any_effect(self) -> None: with self.harness() as h: - h.needs(put=True, get=True, registers=True) - inside = h.endpoint(["A"]) - stray = HostBuffer(SLOT, fill=0x2D) + h.needs(put=True, get=True) h.publish({"A": b"aaaaaaaa"}) - + stray = HostBuffer(SLOT, fill=0x2D) with pytest.raises(DriverRejected): h.get([RowIO(h.key("A"), (stray.span(),))]) assert stray.read() == b"-" * SLOT and stray.canaries_intact() + with pytest.raises(DriverRejected): + h.put([RowIO(h.key("S"), (stray.span(),))]) + reader = h.endpoint(["S"], fill=0x2D) + (result,) = h.settle(h.get(reader.rows())) + h.assert_absent(result, reader, 0) + + def test_a_span_leaving_or_outliving_its_registration_is_rejected(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True, registers=True) + h.publish({"A": b"aaaaaaaa"}) + inside = h.endpoint(["A"], fill=0x2D) + crossing = Span(inside.memory.address + SLOT // 2, SLOT) # starts inside, ends past + with pytest.raises(DriverRejected): + h.get([RowIO(h.key("A"), (crossing,))]) + assert inside.read_row(0) == b"-" * SLOT and inside.intact() inside.fill_row(0, b"bbbbbbbb") inside.unregister() @@ -538,6 +628,18 @@ def test_rows_outside_a_live_registration_are_rejected_before_any_effect(self) - (result,) = h.settle(h.get(reader.rows())) h.assert_absent(result, reader, 0) + def test_a_device_region_is_rejected_unless_device_memory_is_declared(self) -> None: + with self.harness() as h: + h.needs(registers=True) + if h.driver.capabilities.device_memory: + pytest.skip("this driver takes device regions") + memory = HostBuffer(SLOT, fill=0x2D) # stands in for a device region in-process + device = RegionLayout(SLOT, 1, "kv", MemoryKind.DEVICE) + with pytest.raises(DriverRejected): + h.driver.register(memory.address, SLOT, device) + handle = h.driver.register(memory.address, SLOT, RegionLayout(SLOT, 1, "kv")) + h.driver.unregister(handle) + def test_a_key_appears_at_most_once_per_operation(self) -> None: with self.harness() as h: h.needs(put=True, get=True, rows=2) @@ -558,12 +660,14 @@ def test_rows_beyond_the_declared_limit_are_rejected(self) -> None: with self.harness() as h: h.needs(put=True, get=True) limit = h.driver.capabilities.max_rows_per_op - if limit is None or limit >= 3: - pytest.skip("this driver declares no row limit this check can exceed") - writer = h.endpoint(["A", "B", "C"][: limit + 1]) + if limit is None: + pytest.skip("this driver declares no row limit") + if limit >= ROW_LIMIT_CHECKED: + pytest.skip(f"a limit of {limit} rows is above what this check allocates") + writer = h.endpoint([f"R{i}" for i in range(limit + 1)]) with pytest.raises(DriverRejected): h.put(writer.rows()) - reader = h.endpoint(["A"], fill=0x2D) + reader = h.endpoint(["R0"], fill=0x2D) (result,) = h.settle(h.get(reader.rows())) h.assert_absent(result, reader, 0) @@ -584,18 +688,45 @@ def test_an_undeclared_direction_is_rejected(self) -> None: # Results, access and capabilities over time. def test_results_are_stable_and_access_ends_with_evidence(self) -> None: + """Both objects exist: a get of an absent key may conclude FAULT for a driver without + ``absent_is_miss``, and a faulted row may leave access unprovable, which is allowed.""" with self.harness() as h: h.needs(put=True, get=True, rows=2) - h.publish({"A": b"aaaaaaaa"}) - reader = h.endpoint(["A", "C"]) + h.publish({"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + reader = h.endpoint(["A", "B"]) op = h.get(reader.rows()) results = h.settle(op) + assert [r.status for r in results] == [RowStatus.READY, RowStatus.READY] assert h.driver.poll(op) == results h.settle_access(op) assert h.driver.access_ended(op) assert not h.driver.access_unprovable(op) assert h.driver.poll(op) == results + def test_an_unprovable_get_makes_drain_return_at_once_and_keeps_its_region(self) -> None: + staged = self.make_unprovable_get() + if staged is None: + pytest.skip("this driver cannot stage an operation whose end it cannot prove") + driver, lose_the_evidence = staged + with self.harness(driver) as h: + h.needs(get=True) + reader = h.endpoint(["U"], fill=0x2D) + op = h.get(reader.rows()) + lose_the_evidence(op) + h.settle(op) + wait_until( + lambda: driver.access_unprovable(op), "the end to be declared unprovable", h.timeout + ) + assert not driver.access_ended(op) + started = time.monotonic() + assert driver.drain(h.timeout) is False + assert time.monotonic() - started < h.timeout / 2, ( + "drain waited out its budget on an operation that cannot prove its end" + ) + if self.registers_on_demand: + with pytest.raises(DriverRejected): + reader.unregister() + def test_cancel_after_conclusion_changes_nothing(self) -> None: with self.harness() as h: h.needs(put=True, get=True, cancel=True) @@ -607,6 +738,68 @@ def test_cancel_after_conclusion_changes_nothing(self) -> None: assert h.driver.poll(op) == results assert reader.read_row(0) == b"aaaaaaaa" + def test_cancel_of_a_get_inside_the_store_neither_waits_nor_lies(self) -> None: + """Cancel is best effort, so only what it must not do is checked: wait for the store, + call a row the store already took CANCELLED, or end access while that row is inside the + store. The hook guarantees that access cannot end before the gate opens, so the last + holds whatever result the row carries. A driver that hands out the handle only once the + store answered cannot be asked.""" + if not self.handles_before_completion: + pytest.skip("this driver returns the handle of a get only once the store answered") + staged = self.make_blocked_get() + if staged is None: + pytest.skip("this driver cannot stage a get blocked in the store") + driver, blocked_key, blocked, unblock = staged + with self.harness(driver) as h: + h.needs(get=True, cancel=True) + if blocked is None: + pytest.skip("this driver's blocked get gives no signal that it reached the store") + reader = h.endpoint(["S"], fill=0x2D) + handed: dict[str, OpHandle] = {} + handle_back = threading.Event() + returned = threading.Event() + failures: list[BaseException] = [] + + def submit() -> None: + try: + handed["op"] = h.get([RowIO(blocked_key, (reader.span(0),))]) + except BaseException as error: # noqa: BLE001 - reported from the main thread + failures.append(error) + finally: + handle_back.set() + + def cancel() -> None: + driver.cancel(handed["op"]) + returned.set() + + # Both off this thread, so a get or a cancel that waits for the store fails instead + # of hanging; the gate opens on every exit. + submitter = threading.Thread(target=submit, daemon=True) + canceller = threading.Thread(target=cancel, daemon=True) + try: + submitter.start() + assert handle_back.wait(h.timeout), "the get kept its handle" + if failures: + raise AssertionError(f"get raised: {failures[0]!r}") from failures[0] + op = handed["op"] + assert blocked.wait(h.timeout), "the get never reached the store" + canceller.start() + assert returned.wait(h.timeout), "cancel waited for the store" + (during,) = driver.poll(op) + assert during.status is not RowStatus.CANCELLED, ( + "a row the store had taken was called cancelled" + ) + assert not driver.access_ended(op), "access ended with the row in the store" + finally: + unblock() + for thread in (submitter, canceller): + if thread.ident is not None: # started + thread.join(h.timeout) + assert not submitter.is_alive() and not canceller.is_alive() + (after,) = h.settle(op) + if during.concluded: + assert after == during + def test_capabilities_do_not_change(self) -> None: with self.harness() as h: h.needs(put=True) @@ -695,9 +888,12 @@ def test_a_blocked_get_delays_neither_another_get_nor_the_non_blocking_calls(sel seen: dict[str, object] = {} failures: list[BaseException] = [] + handle_back = threading.Event() + def slow_get() -> None: try: - h.get([RowIO(blocked_key, (slow.span(0),))]) + seen["slow_op"] = h.get([RowIO(blocked_key, (slow.span(0),))]) + handle_back.set() except BaseException as error: # noqa: BLE001 - reported from the main thread failures.append(error) @@ -706,6 +902,24 @@ def quick_calls() -> None: op = h.get(quick.rows()) seen["quick"] = h.settle(op) seen["ended"] = h.driver.access_ended(op) + # A driver that returns the handle before the store answers is asked about + # the stalled operation itself; the answers must come at once. A suite that + # declares such a driver waits for the handle, so this branch always runs. + if self.handles_before_completion: + assert handle_back.wait(h.timeout), "the slow get kept its handle" + if handle_back.is_set(): + slow_op = seen["slow_op"] + seen["slow_poll"] = h.driver.poll(slow_op) + seen["slow_ended"] = h.driver.access_ended(slow_op) + seen["slow_unprovable"] = h.driver.access_unprovable(slow_op) + seen["slow_checked"] = True + if self.registers_on_demand: + try: + slow.unregister() + except DriverRejected: + seen["unregister"] = "refused" + else: + seen["unregister"] = "accepted" seen["drain"] = h.driver.drain(0.05) except BaseException as error: # noqa: BLE001 - reported from the main thread failures.append(error) @@ -731,6 +945,15 @@ def quick_calls() -> None: assert [r.status for r in quick_results] == [RowStatus.READY] assert quick.read_row(0) == b"bbbbbbbb" assert seen["drain"] is False # the slow get is still inside the store + if self.registers_on_demand: + assert seen["unregister"] == "refused", "a busy region was unregistered" + if self.handles_before_completion: + assert seen.get("slow_checked"), "the stalled get was never asked about" + if seen.get("slow_checked"): + assert seen["slow_ended"] is False and seen["slow_unprovable"] is False + slow_results = seen["slow_poll"] + assert isinstance(slow_results, tuple) + assert all(r.status is not RowStatus.READY for r in slow_results) finally: unblock() slow_thread.join(h.timeout) @@ -799,14 +1022,24 @@ def test_probe_raises_on_a_store_error_rather_than_answering(self) -> None: # Registration, drain and close. - def test_unregister_is_idempotent_and_frees_the_range(self) -> None: + def test_unregister_is_idempotent(self) -> None: + """Also for a driver whose memory was registered up front: its one registration is + unregistered twice, and the suite's shutdown hook must cope with that.""" + with self.harness() as h: + endpoint = h.endpoint(["A"]) + assert endpoint.handle is not None + assert h.driver.drain(self.timeout) is True + h.driver.unregister(endpoint.handle) + h.driver.unregister(endpoint.handle) + endpoint.released = True # the harness must not give pool memory back twice + + def test_unregister_frees_the_range_for_a_new_registration(self) -> None: with self.harness() as h: h.needs(registers=True) endpoint = h.endpoint(["A"]) with pytest.raises(DriverRejected): h.driver.register(endpoint.memory.address, SLOT, RegionLayout(SLOT, 1)) endpoint.unregister() - endpoint.unregister() try: handle = h.driver.register(endpoint.memory.address, SLOT, RegionLayout(SLOT, 1)) except DriverRejected as error: @@ -826,17 +1059,22 @@ def test_drain_is_one_way(self) -> None: assert writer.read_row(0) == b"aaaaaaaa" def test_close_needs_a_true_drain_and_no_live_registrations(self) -> None: + """For a driver whose memory was registered up front, the suite's shutdown hook does + the unregistering; the order it has to keep is the same.""" with self.harness() as h: - h.needs(put=True, registers=True) + h.needs(put=True) writer = h.publish({"A": b"aaaaaaaa"}) with pytest.raises(DriverRejected): h.driver.close() assert h.driver.drain(self.timeout) is True with pytest.raises(DriverRejected): h.driver.close() - writer.unregister() - h.driver.close() - h.driver.close() + if self.registers_on_demand: + writer.unregister() + h.driver.close() + else: + self.shutdown(h.driver) + h.driver.close() # idempotent h.done = True # closed by the check itself with pytest.raises(DriverRejected): h.put(writer.rows()) diff --git a/tensorrt_llm/_torch/disaggregation/testing/memory_store.py b/tensorrt_llm/_torch/disaggregation/testing/memory_store.py index 746be8114816..88471bdad7df 100644 --- a/tensorrt_llm/_torch/disaggregation/testing/memory_store.py +++ b/tensorrt_llm/_torch/disaggregation/testing/memory_store.py @@ -28,6 +28,7 @@ from dataclasses import dataclass from typing import Literal +from tensorrt_llm._torch.disaggregation.base.memory import MemoryKind from tensorrt_llm._torch.disaggregation.blob.store import ( DriverCapabilities, DriverRejected, @@ -54,6 +55,7 @@ probe=True, max_rows_per_op=None, absent_is_miss=True, + device_memory=True, # a DEVICE region is plain memory in-process; the plumbing is what counts multi_span=True, ) MAX_PROBES_IN_FLIGHT = 4 @@ -213,6 +215,12 @@ class _Probe: finished: threading.Event +@dataclass(eq=False) +class _Worker: + thread: threading.Thread + finished: threading.Event + + def _gather(spans: Sequence[Span]) -> bytes: return b"".join(ctypes.string_at(span.address, span.nbytes) for span in spans) @@ -232,10 +240,15 @@ def __init__( store: MemoryStore, *, controlled: bool = False, + background: bool = False, capabilities: DriverCapabilities = DEFAULT_CAPABILITIES, ) -> None: + if controlled and background: + raise ValueError("a driver is controlled or runs in the background, not both") self.store = store self.controlled = controlled + self.background = background # put/get return once recorded; a thread executes them + self._workers: list[_Worker] = [] self._capabilities = capabilities self._lock = threading.Condition(threading.RLock()) self._registrations: list[_Registration] = [] @@ -263,6 +276,7 @@ def register(self, address: int, nbytes: int, layout: RegionLayout) -> _Registra raise DriverRejected(str(error)) from None if layout.nbytes > nbytes: raise DriverRejected(f"layout needs {layout.nbytes} bytes, region has {nbytes}") + self._check_region_kind(layout) with self._lock: self._require_accepting() for other in self._registrations: @@ -369,8 +383,11 @@ def close(self) -> None: self._check_closable() self._closed = True probes, self._probes = self._probes, [] + workers, self._workers = self._workers, [] for probe in probes: # all finished; join only lets the threads exit probe.worker.join() + for worker in workers: # all finished too; see _quiet + worker.thread.join() # Controlled mode: the test releases each stage. @@ -527,6 +544,10 @@ def _check_disjoint(self, kind: Kind, rows: Sequence[RowIO]) -> None: if first.overlaps(second): raise DriverRejected(f"spans {first} and {second} overlap within one get") + def _check_region_kind(self, layout: RegionLayout) -> None: + if layout.memory_kind is MemoryKind.DEVICE and not self._capabilities.device_memory: + raise DriverRejected("this driver registers host memory only") + def _check_registered(self, rows: Sequence[RowIO]) -> None: for io in rows: for span in io.spans: @@ -564,9 +585,12 @@ def _drain_can_return(self) -> bool: return self._quiet() or any(op.unprovable and not op.access_ended for op in self._ops) def _quiet(self) -> bool: - """Nothing can still touch memory and no background work is running.""" - return all(op.access_ended for op in self._ops) and all( - probe.finished.is_set() for probe in self._probes + """Nothing can still touch memory and no background work is running. An operation's + access ending is not enough: its worker may still be on its way out.""" + return ( + all(op.access_ended for op in self._ops) + and all(probe.finished.is_set() for probe in self._probes) + and all(worker.finished.is_set() for worker in self._workers) ) # Internals. @@ -586,10 +610,34 @@ def _submit(self, kind: Kind, rows: Sequence[RowIO], hint: Hint = None) -> Memor self._ops.append(op) if hint is not None: self.hints_seen.append(hint) - if not self.controlled: + if self.background: + # Recorded, started and registered under one lock, so a drain or close that + # follows sees a worker that exists and has started. + finished = threading.Event() + thread = threading.Thread( + target=self._run, args=(op, finished), name="memory-driver-op", daemon=True + ) + try: + thread.start() + except RuntimeError as error: + # The thread never ran, so no row started: the recorded operation concludes + # as faults and, having touched nothing, has ended its access. + self._fail_unstarted(op, f"no worker thread: {error}") + else: + self._workers.append(_Worker(thread, finished)) + if not self.background and not self.controlled: + self._run(op, None) + return op + + def _run(self, op: MemoryOp, finished: threading.Event | None) -> None: + try: self._before_execute(op) self._execute(op) - return op + finally: + if finished is not None: + finished.set() + with self._lock: + self._lock.notify_all() def _before_execute(self, op: MemoryOp) -> None: """Between recording an operation and executing it; a test races a drain here.""" @@ -655,6 +703,13 @@ def _verify(self, io: RowIO, stored: StoredObject) -> str | None: """Why the stored object may not be delivered as ``io``, or ``None``.""" return check_trailer(stored.trailer, io.key.encode(), io.nbytes, stored.payload) + def _fail_unstarted(self, op: MemoryOp, reason: str) -> None: + """Every row FAULT and access over, for an operation none of whose rows started.""" + for row in op.rows: + row.published = faulted(reason) + op.access_ended = True + self._lock.notify_all() + def _cancel_unstarted(self, op: MemoryOp) -> None: for row in op.rows: if not row.started and not row.published.concluded: diff --git a/tests/unittest/_torch/disaggregation/blob/test_blob_format.py b/tests/unittest/_torch/disaggregation/blob/test_blob_format.py index 3c88e0b6e147..e3a4ce336e14 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_blob_format.py +++ b/tests/unittest/_torch/disaggregation/blob/test_blob_format.py @@ -6,6 +6,7 @@ import pytest +from tensorrt_llm._torch.disaggregation.base.memory import MemoryKind from tensorrt_llm._torch.disaggregation.blob.keys import FORMAT_VERSION, encode_key from tensorrt_llm._torch.disaggregation.blob.store import ( DriverCapabilities, @@ -82,6 +83,8 @@ def test_span_and_layout_validate_their_numbers() -> None: with pytest.raises(ValueError): RegionLayout(slot_bytes, slot_count) assert RegionLayout(8, 4).nbytes == 32 + assert RegionLayout(8, 4).memory_kind is MemoryKind.HOST + assert RegionLayout(8, 4, "kv", MemoryKind.DEVICE).memory_kind is MemoryKind.DEVICE def test_a_row_is_a_key_and_disjoint_spans_in_the_order_given() -> None: diff --git a/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py index c8ed766c4496..98e691553725 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py +++ b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py @@ -13,6 +13,7 @@ import pytest +from tensorrt_llm._torch.disaggregation.base.memory import MemoryKind from tensorrt_llm._torch.disaggregation.blob.store import ( DriverCapabilities, DriverRejected, @@ -22,7 +23,6 @@ RegionLayout, RegistrationError, RowIO, - RowResult, RowStatus, Span, StoreDriver, @@ -30,7 +30,10 @@ concluded, faulted, ) -from tensorrt_llm._torch.disaggregation.testing.broken_drivers import BROKEN_CONFORMANCE +from tensorrt_llm._torch.disaggregation.testing.broken_drivers import ( + BROKEN_ASYNC_CONFORMANCE, + BROKEN_CONFORMANCE, +) from tensorrt_llm._torch.disaggregation.testing.driver_conformance import ( SLOT, UNRELEASED_ENDPOINTS, @@ -41,11 +44,13 @@ Unblock, key_of, statuses, + wait_until, ) from tensorrt_llm._torch.disaggregation.testing.host_memory import GUARD, HostBuffer, MemoryWindow from tensorrt_llm._torch.disaggregation.testing.memory_store import ( DEFAULT_CAPABILITIES, MAX_PROBES_IN_FLIGHT, + Commit, Kind, MemoryDriver, MemoryOp, @@ -73,7 +78,7 @@ def __init__(self, fail_on_commit: int) -> None: self._fail_on = fail_on_commit self._commits = 0 - def commit(self, key: str, payload: bytes, trailer: bytes): # type: ignore[no-untyped-def] + def commit(self, key: str, payload: bytes, trailer: bytes) -> Commit: self._commits += 1 if self._commits == self._fail_on: raise StoreError("connection dropped") @@ -96,7 +101,7 @@ def get(self, key: str) -> StoredObject | None: class CrashingStore(MemoryStore): """Raises something that is not a ``StoreError`` from inside a commit, as a binding bug would.""" - def commit(self, key: str, payload: bytes, trailer: bytes): # type: ignore[no-untyped-def] + def commit(self, key: str, payload: bytes, trailer: bytes) -> Commit: raise RuntimeError("binding returned garbage") @@ -116,7 +121,7 @@ def _hold(self, key: str) -> None: self.waiting.set() self.gate.wait() - def commit(self, key: str, payload: bytes, trailer: bytes): # type: ignore[no-untyped-def] + def commit(self, key: str, payload: bytes, trailer: bytes) -> Commit: self._hold(key) return super().commit(key, payload, trailer) @@ -154,11 +159,14 @@ class TestMemoryDriverConformance(DriverConformance): driver_cls: type[MemoryDriver] = MemoryDriver capabilities = DEFAULT_CAPABILITIES + background = False # True: put/get return once recorded and a thread executes them def _driver(self, store: MemoryStore | None = None) -> MemoryDriver: # An empty store is falsy (it has a length), so test for None explicitly. return self.driver_cls( - store if store is not None else MemoryStore(), capabilities=self.capabilities + store if store is not None else MemoryStore(), + background=self.background, + capabilities=self.capabilities, ) def make_driver(self) -> StoreDriver: @@ -190,6 +198,20 @@ def make_failing_presence(self) -> tuple[StoreDriver, str]: store.fail_next_presence("connection dropped") return self._driver(store), STAGED_KEY + def make_failing_get(self) -> tuple[StoreDriver, str]: + return self._driver(TimeoutStore(STAGED_KEY)), STAGED_KEY + + def make_unprovable_get(self) -> tuple[StoreDriver, Callable[[OpHandle], None]]: + driver = self.driver_cls(MemoryStore(), controlled=True, capabilities=self.capabilities) + + def lose_the_evidence(op: OpHandle) -> None: + assert isinstance(op, MemoryOp) + driver.transfer(op) + driver.publish_results(op) + driver.mark_unprovable(op) + + return driver, lose_the_evidence + class _HintFreeSignatures(MemoryDriver): """A driver written before hints existed: ``get`` and ``probe`` take no hint. With @@ -206,6 +228,13 @@ class TestHintFreeDriverConformance(TestMemoryDriverConformance): driver_cls = _HintFreeSignatures +class TestAsyncMemoryDriverConformance(TestMemoryDriverConformance): + """The same driver returning handles before the store answers, as a real SDK does.""" + + background = True + handles_before_completion = True + + SHARED_STORE = MemoryStore() @@ -234,8 +263,8 @@ def make_driver(self) -> StoreDriver: self.next_offset = 0 return driver - def allocate(self, nbytes: int, fill: int) -> Allocation: - start = self.next_offset + GUARD + def allocate(self, nbytes: int, fill: int, memory_kind: MemoryKind) -> Allocation: + start = self.next_offset + GUARD # host memory stands in for either kind in-process self.next_offset = start + nbytes + GUARD assert self.next_offset <= self.POOL_BYTES, "the pool is too small for this check" return MemoryWindow(self.pool.address + start, nbytes, fill, owner=self.pool), self.handle @@ -686,6 +715,34 @@ def test_a_drain_racing_a_submission_leaves_a_cancelled_operation_with_its_handl assert len(store) == 0 assert driver.drain(TIMEOUT) is True + def test_a_worker_that_cannot_start_leaves_a_faulted_operation_with_its_handle(self) -> None: + """Background mode records the operation before its thread starts. When the thread cannot + start, nothing ran: every row is a fault, access is over, and the driver goes on.""" + store = MemoryStore() + driver = self.make(store, background=True) + writer = endpoint(driver, ["A", "B"]) + writer.fill_row(0, b"aaaaaaaa") + writer.fill_row(1, b"bbbbbbbb") + + def refuse(thread: threading.Thread) -> None: + raise RuntimeError("can't start new thread") + + with pytest.MonkeyPatch.context() as patched: + patched.setattr(threading.Thread, "start", refuse) + try: + op = driver.put(writer.rows()) + except Exception as error: # noqa: BLE001 - the check is that the handle comes back + pytest.fail(f"put raised instead of returning its handle: {error!r}") + results = driver.poll(op) + assert [r.status for r in results] == [RowStatus.FAULT, RowStatus.FAULT] + assert results[0].fault == "no worker thread: can't start new thread" + assert driver.access_ended(op) + assert len(store) == 0 + later = driver.put(writer.rows()) # nothing stale stands in the way of the next thread + wait_until(lambda: concluded(driver.poll(later)), "the next operation", TIMEOUT) + assert statuses(driver, later) == [RowStatus.READY, RowStatus.READY] + self._shutdown(driver) + def test_unregister_is_refused_while_an_operation_may_access_the_region(self) -> None: driver = self.make(controlled=True) reader = endpoint(driver, ["A"]) @@ -934,19 +991,6 @@ def test_harness_retains_endpoints_of_an_unprovable_operation_without_complaint( # Drivers that break one rule only the in-memory tests can observe. -class _UnregistersBusyRegions(MemoryDriver): - def _check_region_idle(self, registration: object) -> None: - pass - - -class _CallsATimeoutAMiss(MemoryDriver): - def _get_row(self, io: RowIO) -> RowResult: - try: - return super()._get_row(io) - except StoreError: - return RowResult(RowStatus.MISS) - - class _SkipsTheIntegrityCheck(MemoryDriver): def _verify(self, io: RowIO, stored: StoredObject) -> str | None: return None @@ -984,21 +1028,13 @@ def _execute(self, op: MemoryOp) -> None: self.end_access(op) -class _DrainWaitsForEveryone(MemoryDriver): - def _drain_can_return(self) -> bool: - return self._quiet() +class _LeavesAnUnstartedOperationPending(MemoryDriver): + def _fail_unstarted(self, op: MemoryOp, reason: str) -> None: + pass # rule -> (driver breaking it, the memory-driver check that must fail) BROKEN_MEMORY: dict[str, tuple[type[MemoryDriver], str]] = { - "unregister needs an idle region": ( - _UnregistersBusyRegions, - "test_unregister_is_refused_while_an_operation_may_access_the_region", - ), - "a store timeout is a fault": ( - _CallsATimeoutAMiss, - "test_a_store_timeout_is_a_fault_not_a_miss", - ), "integrity is verified on every get": ( _SkipsTheIntegrityCheck, "test_a_corrupted_object_is_a_fault_and_never_lands", @@ -1027,9 +1063,9 @@ def _drain_can_return(self) -> bool: _ExecutesWithTheControls, "test_a_drain_racing_a_submission_leaves_a_cancelled_operation_with_its_handle", ), - "drain returns at once on an unprovable operation": ( - _DrainWaitsForEveryone, - "test_drain_returns_false_at_once_even_while_another_operation_is_still_provable", + "a worker that cannot start concludes its operation": ( + _LeavesAnUnstartedOperationPending, + "test_a_worker_that_cannot_start_leaves_a_faulted_operation_with_its_handle", ), } @@ -1048,12 +1084,12 @@ def _run_check(suite: object, check: str) -> None: raise UnexpectedSkip(f"{check} skipped instead of running: {skipped}") from None -def _conformance(driver_cls: type[MemoryDriver], caps: DriverCapabilities) -> DriverConformance: - suite = type( - "Suite", - (TestMemoryDriverConformance,), - {"driver_cls": driver_cls, "capabilities": caps, "timeout": 0.3}, - ) +def _conformance( + driver_cls: type[MemoryDriver], + caps: DriverCapabilities, + base: type[TestMemoryDriverConformance] = TestMemoryDriverConformance, +) -> DriverConformance: + suite = type("Suite", (base,), {"driver_cls": driver_cls, "capabilities": caps, "timeout": 0.3}) return suite() # type: ignore[no-any-return] @@ -1067,6 +1103,17 @@ def test_each_conformance_check_passes_the_driver_and_fails_the_one_breaking_its _run_check(_conformance(broken, capabilities), check) +@pytest.mark.parametrize("rule", sorted(BROKEN_ASYNC_CONFORMANCE)) +def test_each_async_conformance_check_passes_the_driver_and_fails_the_one_breaking_its_rule( + rule: str, +) -> None: + broken, capabilities, check = BROKEN_ASYNC_CONFORMANCE[rule] + asynchronous = TestAsyncMemoryDriverConformance + _run_check(_conformance(MemoryDriver, capabilities, asynchronous), check) + with pytest.raises(CHECK_FAILED): + _run_check(_conformance(broken, capabilities, asynchronous), check) + + @pytest.mark.parametrize("rule", sorted(BROKEN_MEMORY)) def test_each_memory_driver_check_fails_the_driver_breaking_its_rule(rule: str) -> None: broken, check = BROKEN_MEMORY[rule] @@ -1086,10 +1133,11 @@ def test_a_skip_on_the_broken_side_is_not_an_expected_failure() -> None: def test_each_broken_driver_breaks_one_rule_and_names_an_existing_check() -> None: - conformance = [cls for cls, _, _ in BROKEN_CONFORMANCE.values()] + pairings = [*BROKEN_CONFORMANCE.values(), *BROKEN_ASYNC_CONFORMANCE.values()] + conformance = [cls for cls, _, _ in pairings] memory = [cls for cls, _ in BROKEN_MEMORY.values()] assert len(set(conformance + memory)) == len(conformance) + len(memory) - checks = {check for _, _, check in BROKEN_CONFORMANCE.values()} + checks = {check for _, _, check in pairings} assert checks <= {name for name in dir(DriverConformance) if name.startswith("test_")} checks = {check for _, check in BROKEN_MEMORY.values()} assert checks <= {name for name in dir(TestMemoryDriver) if name.startswith("test_")} diff --git a/tests/unittest/_torch/disaggregation/blob/test_testing_isolation.py b/tests/unittest/_torch/disaggregation/blob/test_testing_isolation.py index 92eb578ea560..87ae29c02566 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_testing_isolation.py +++ b/tests/unittest/_torch/disaggregation/blob/test_testing_isolation.py @@ -21,12 +21,7 @@ def _package_of(path: pathlib.Path) -> str: """The package a module at ``path`` resolves relative imports against.""" relative = path.resolve().relative_to(PACKAGE_DIR.parent).with_suffix("") - parts = list(relative.parts) - if parts[-1] != "__init__": - parts.pop() - else: - parts.pop() - return ".".join(parts) if path.name == "__init__.py" else ".".join(parts) + return ".".join(relative.parts[:-1]) # a module's parent; a package's own name def _testing_imports(source: str, package: str) -> list[str]: