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 new file mode 100644 index 000000000000..efab37bb6b99 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/blob/__init__.py @@ -0,0 +1,8 @@ +# 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. + +So far this package holds the private driver boundary (``store.py``) and the key encoding +(``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/keys.py b/tensorrt_llm/_torch/disaggregation/blob/keys.py new file mode 100644 index 000000000000..a27af1c7d481 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/blob/keys.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Store keys for blob objects. + +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 + +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/"``. + + Args: + 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.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{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..acf2d60fede9 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/blob/store.py @@ -0,0 +1,410 @@ +# 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 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/``, from which only +``MemoryKind`` is shared; the backend translates between the two. +""" + +from __future__ import annotations + +import enum +from collections.abc import Sequence +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 +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 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. 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. + + Args: + reason: What went wrong. + handle: The handle to ``unregister`` later when the store may have registered the region + anyway, or ``None`` when the store certainly did not. + """ + + 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. + + 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" + """Not concluded yet.""" + + READY = "ready" + """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 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 + 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; 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. Says nothing + about the destination's contents.""" + + +@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 + + +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: + """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 + + 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. + ``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: + 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 and the spans its payload occupies. + + 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 + 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): + """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" + NOT_READY = "not_ready" + UNKNOWN = "unknown" + + +@dataclass(frozen=True) +class DriverCapabilities: + """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. + 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. + 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 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. + """ + + get: bool + put: bool + cancel: bool + 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): + 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``, ``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. + + 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. + + 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 + def capabilities(self) -> DriverCapabilities: + """The same value on every call.""" + ... + + def register(self, address: int, nbytes: int, layout: RegionLayout) -> RegionHandle: + """Register one region the driver may read and write. + + ``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, 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. + """ + ... + + def unregister(self, handle: RegionHandle) -> None: + """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. + + The caller guarantees that the rows' spans hold the payload and stay unchanged until + ``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``, 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. + + 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``. + """ + ... + + def poll(self, op: OpHandle) -> tuple[RowResult, ...]: + """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``. 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: + """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, 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, and with a hint only when ``capabilities.hints`` is. + 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 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/tensorrt_llm/_torch/disaggregation/testing/__init__.py b/tensorrt_llm/_torch/disaggregation/testing/__init__.py new file mode 100644 index 000000000000..6ad7306bcf28 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/testing/__init__.py @@ -0,0 +1,22 @@ +# 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. +""" + +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 new file mode 100644 index 000000000000..69235c6e2b1d --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py @@ -0,0 +1,469 @@ +# 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 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, MemoryStore, _scatter + +_MISS = RowResult(RowStatus.MISS) +_READY = RowResult(RowStatus.READY) +_CANCELLED = RowResult(RowStatus.CANCELLED) + + +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 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) + 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 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) + 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 RegistersAnyMemoryKind(MemoryDriver): + def _check_region_kind(self, layout: RegionLayout) -> 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, + device_memory=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 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) + 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 DrainWaitsForEveryone(MemoryDriver): + def _drain_can_return(self) -> bool: + return self._quiet() + + +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 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]] = { + "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_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, + 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 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, + "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", + ), + "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, + "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", + ), + "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, + "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 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, + "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_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 new file mode 100644 index 000000000000..a5d672296024 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py @@ -0,0 +1,1080 @@ +# 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.base.memory import MemoryKind +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 +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. +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, + 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. + 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", self.memory_kind) + 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, + 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() + 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, 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. + + 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, 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 + + 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. + + 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, 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. 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 + + 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. 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: + """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 + + 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 + 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" + + 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 + 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() + + 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_a_row_outside_every_registration_is_rejected_before_any_effect(self) -> None: + with self.harness() as h: + 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() + 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_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) + 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: + 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(["R0"], 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: + """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", "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) + 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_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) + 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] = [] + + handle_back = threading.Event() + + def slow_get() -> None: + try: + 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) + + def quick_calls() -> None: + try: + 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) + + # 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 + 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) + 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(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() + 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: + """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) + 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() + 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/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/tensorrt_llm/_torch/disaggregation/testing/memory_store.py b/tensorrt_llm/_torch/disaggregation/testing/memory_store.py new file mode 100644 index 000000000000..88471bdad7df --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/testing/memory_store.py @@ -0,0 +1,757 @@ +# 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 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, 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 +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.base.memory import MemoryKind +from tensorrt_llm._torch.disaggregation.blob.store import ( + DriverCapabilities, + DriverRejected, + Hint, + OpHandle, + ProbeStatus, + RegionHandle, + RegionLayout, + RegistrationError, + RowIO, + RowResult, + RowStatus, + Span, + StoreError, + 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, + 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 + +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.""" + + +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. + + 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: + 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 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 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 + + def contains(self, key: str) -> bool: + return self.presence(key) is ProbeStatus.HIT + + 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 stored + + 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 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: + 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: + 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: Kind + rows: list[_Row] + hint: Hint = None + access_ended: bool = False + unprovable: bool = False + + +@dataclass(eq=False) +class _Probe: + worker: threading.Thread + 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) + + +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.""" + + def __init__( + self, + 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] = [] + 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 + 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. + + @property + def capabilities(self) -> DriverCapabilities: + return self._capabilities + + def register(self, address: int, nbytes: int, layout: RegionLayout) -> _Registration: + try: + span = Span(address, nbytes) + except ValueError as error: + raise DriverRejected(str(error)) from None + if layout.nbytes > nbytes: + raise DriverRejected(f"layout needs {layout.nbytes} bytes, region has {nbytes}") + self._check_region_kind(layout) + with self._lock: + self._require_accepting() + for other in self._registrations: + if other.live and other.span.overlaps(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 + 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) + 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) + registration.live = False + + def put(self, rows: Sequence[RowIO]) -> MemoryOp: + return self._submit("put", 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: + 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 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, 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() + + 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._probes = [probe for probe in self._probes if not probe.finished.is_set()] + 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 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] + return tuple(answers.get(key, ProbeStatus.UNKNOWN) for key in keys) + + def drain(self, budget_s: float) -> bool: + with self._lock: + self._require_not_closed() + self._start_draining() + for op in self._ops: + self._cancel_unstarted(op) + 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, [] + 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. + + 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; any raise 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, 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): + 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 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: + 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") + 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.""" + 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 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 + + 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] + + 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 ever started to exit, for tests that opened a gate they + had closed.""" + with self._lock: + workers = list(self._probe_workers) + for worker in workers: + 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 _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") + limit = self._capabilities.max_rows_per_op + if limit is not None and len(rows) > limit: + 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_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_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: + 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(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") + + 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 _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 + + 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. 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. + + 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_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) + if hint is not None: + self.hints_seen.append(hint) + 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) + 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.""" + + 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: + 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 + 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: + try: + stored = self.store.get(io.key) + except ObjectNotReady: + return _NOT_READY + if stored is None: + 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) + _scatter(io.spans, stored.payload) + return _READY + + def _trailer_for(self, io: RowIO, payload: bytes) -> bytes: + 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.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: + 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") + + @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/tensorrt_llm/_torch/disaggregation/testing/object_trailer.py b/tensorrt_llm/_torch/disaggregation/testing/object_trailer.py new file mode 100644 index 000000000000..2da2d42e679f --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/testing/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..e3a4ce336e14 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_blob_format.py @@ -0,0 +1,174 @@ +# 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 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, + 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 + + +# 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_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 + with pytest.raises(ValueError): + RowResult(RowStatus.FAULT) + with pytest.raises(ValueError): + RowResult(RowStatus.READY, "text") + + +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 + 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: + 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("k", ()) + with pytest.raises(ValueError): + RowIO("k", (Span(0x1000, 64), Span(0x1020, 64))) + + +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): + 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) + 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] + + +# 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..98e691553725 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py @@ -0,0 +1,1143 @@ +# 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 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 gc +import threading +import time +import weakref +from collections.abc import Callable, Sequence + +import pytest + +from tensorrt_llm._torch.disaggregation.base.memory import MemoryKind +from tensorrt_llm._torch.disaggregation.blob.store import ( + DriverCapabilities, + DriverRejected, + Hint, + OpHandle, + ProbeStatus, + RegionLayout, + RegistrationError, + RowIO, + RowStatus, + Span, + StoreDriver, + StoreError, + concluded, + faulted, +) +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, + Allocation, + DriverConformance, + Endpoint, + Harness, + 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, + MemoryStore, + StoredObject, +) +from tensorrt_llm._torch.disaggregation.testing.object_trailer import encode_trailer + +pytestmark = pytest.mark.cpu_only + +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 + + +# 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) -> Commit: + 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) -> Commit: + raise RuntimeError("binding returned garbage") + + +class GatedStore(MemoryStore): + """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__() + 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) -> Commit: + 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 + + +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 + 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(), + background=self.background, + 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 + + 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 + ``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 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() + + +class TestSharedStoreConformance(DriverConformance): + """Fresh clients of one persistent store: every check finds the previous checks' objects.""" + + def make_driver(self) -> StoreDriver: + return MemoryDriver(SHARED_STORE) + + +class TestPoolDriverConformance(DriverConformance): + """A driver that registered its memory once, when built; endpoints are carved out of it.""" + + 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, 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 + + def shutdown(self, driver: StoreDriver) -> None: + driver.unregister(self.handle) + driver.close() + + +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_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.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) + + # 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" + + # 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" + + # 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_a_blocked_probe_holds_drain_until_the_store_answers(self) -> None: + store = GatedStore(held={KEY_A}) + 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: + 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() + driver.join_probes(TIMEOUT) + assert driver.probe([KEY_A], TIMEOUT) == (ProbeStatus.UNKNOWN,) # 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 + + # 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) + 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_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", "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, + ) -> 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_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"]) + 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 + + # Registration and shutdown. + + 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.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. + + 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() == [] + + 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() == [] + + 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 + + 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() + + +# Drivers that break one rule only the in-memory tests can observe. + + +class _SkipsTheIntegrityCheck(MemoryDriver): + def _verify(self, io: RowIO, stored: StoredObject) -> str | None: + return None + + +class _ProbesWithoutLimit(MemoryDriver): + def _probe_slot_free(self) -> bool: + return True # the workers stay on the books, so they are still joined + + +class _ReportsAccessEndedEarly(MemoryDriver): + def access_ended(self, op: OpHandle) -> bool: + return True + + +class _AcceptsHintsItCannotRead(MemoryDriver): + 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) + self.publish_results(op) + self.end_access(op) + + +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]] = { + "integrity is verified on every get": ( + _SkipsTheIntegrityCheck, + "test_a_corrupted_object_is_a_fault_and_never_lands", + ), + "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", + ), + "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", + ), + "a worker that cannot start concludes its operation": ( + _LeavesAnUnstartedOperationPending, + "test_a_worker_that_cannot_start_leaves_a_faulted_operation_with_its_handle", + ), +} + +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, + 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] + + +@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] + _run_check(_conformance(MemoryDriver, capabilities), check) + with pytest.raises(CHECK_FAILED): + _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] + suite = type("Suite", (TestMemoryDriver,), {"driver_cls": broken})() + with pytest.raises(CHECK_FAILED): + _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: + 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 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 new file mode 100644 index 000000000000..87ae29c02566 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_testing_isolation.py @@ -0,0 +1,78 @@ +# 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("") + return ".".join(relative.parts[:-1]) # a module's parent; a package's own name + + +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" + )