diff --git a/tensorrt_llm/_torch/disaggregation/blob/README.md b/tensorrt_llm/_torch/disaggregation/blob/README.md new file mode 100644 index 000000000000..223f1a0bc9c9 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/blob/README.md @@ -0,0 +1,225 @@ + + + +# The blob backend + +A `Backend` of the KV sharing contract (`tensorrt_llm/_torch/disaggregation/base/backend.py`) +that moves KV rows between the lender's registered memory and a byte store, through a private +store driver. It is the reference path for any store that can hold a row as an object under a +key: the in-memory store the tests use, Mooncake, KVCR, or another SDK. The backend knows no +store; everything store-specific lives in a driver. + +## Where it sits + +```text +KV cache manager lends regions; keeps the region hold + │ + ▼ +lease holder (runtime) submits fetch / publish views; decides from the outcome; returns the + │ lease once the attempt quiesced + ▼ +BlobBackend this package: checks a view, cuts it into calls and hands each to the + │ driver as one RowBatch; records what the driver tells of each call + ▼ +StoreDriver one per store SDK: registers memory, runs the store's work on threads of + │ its own, moves whole rows, answers two questions apart (result of a row, + ▼ end of access to its memory) and tells of each change through on_change +store SDK +``` + +| Module | Holds | +|---|---| +| `backend.py` | `BlobBackend`, `BlobAttempt`, `BlobConfig`, `DriverCallLimit`, `RegionBinding`, `ConstructionRejected` | +| `store.py` | the driver boundary: `StoreDriver`, the optional `RegistersRegionsTogether`, `RowBatch`, `RowIO`, `OnChange`, `RowStatus`, `DriverCapabilities`, `DriverRejected`, `RegistrationError` | +| `keys.py` | the store key encoding | +| `../testing/` | what a driver's own CI imports from the wheel: the conformance suite, the in-memory store and drivers, the broken drivers that prove every check bites | + +## Running the tests + +```bash +pytest tests/unittest/_torch/disaggregation/blob +``` + +All of it is CPU-only. The files, top down: + +- `test_blob_example.py`: the backend end to end, each test the smallest example of one part + of the picture: two endpoints sharing objects, a lease holder's decisions and leases, the + region hold at shutdown, two backends in one process. +- `test_blob_backend.py`: the backend's behaviour against the in-memory driver, in every + mode the test fixtures can stage (blocking, asynchronous, controlled step by step). +- `test_memory_driver.py`, `test_together_driver.py`: the conformance suite over the in-memory + drivers, and the pairing meta-test that runs each check against a driver breaking exactly the + rule it guards. +- `test_blob_format.py`: keys, spans, layouts, row batches, the object trailer. +- `test_testing_isolation.py`: no engine module imports the test tooling. + +## Writing a driver + +The boundary is `store.py`; its docstrings are the specification. The rules that shape a driver: + +1. **`DriverRejected` is the only raise that certifies nothing escaped.** It is raised before the + store or the caller's memory was touched. A `put` or `get` that raises anything else once + it was made may have touched its rows: the backend keeps them protected for good and sends + the driver no more work. A `register` that fails says through `RegistrationError` what it + left behind: `handle=None` certifies nothing, a handle names what to revoke; a failure the + driver has no handle for escapes as the exception it is, and the memory stays protected. +2. **Two questions, answered apart.** `poll` gives each row's logical result; `access_ended` + says whether the operation can still touch the rows' memory. Neither implies the other. A + row's result may land long before its access ended, and the memory is reused only on the + latter. `access_unprovable` is the driver's admission that it will never know. The driver + tells the caller of each change through the `on_change` it was handed, after the change + shows through those queries; the backend never polls. +3. **The driver owns the store's work.** `put` and `get` take a `RowBatch`, record and queue + it, and return without waiting on the store; the driver's own threads move the bytes, in the + order and at the concurrency its store wants. A batch's rows are laid out alike and never + name a key twice, and every span lies inside one registered extent, so a driver checks memory + once per extent. A driver that cannot take a call refuses it with `DriverRejected`. +4. **Only READY and FAULT are verdicts.** MISS, NOT_READY, DECLINED and CANCELLED say that + nothing was delivered and nothing went wrong, each under a condition the store confirmed; + an error the driver cannot place is FAULT. +5. **`drain` is one-way and never forces an end.** After it the driver accepts no new work; + it returns True only when nothing can still touch memory. `close` is allowed only after a + True drain and with every registration revoked. +6. **Registration comes in two shapes.** A store that takes memory a region at a time + implements `register` and `unregister`. A store that takes all of its memory in one step, + as one that opens its client with it, implements `RegistersRegionsTogether` instead: + `register_regions(specs)` returns one handle for the batch and `unregister_regions(handle)` + gives it back. No atomicity is promised; the failure shapes are the same as `register`'s. + +### The conformance suite + +```python +from tensorrt_llm._torch.disaggregation.testing.driver_conformance import DriverConformance + + +class TestMyDriverConformance(DriverConformance): + def make_driver(self): + return MyDriver(open_my_client()) +``` + +- **Scenario hooks** stage what a check needs from the store: `make_blocked_get`, + `make_held_key`, `make_full_store`, `make_failing_get`, `make_unprovable_get`, and the rest + listed on `DriverConformance`. A hook that returns `None` skips its check; a driver's CI + should implement every hook its store can honour. +- **Pool mode**: a driver whose memory is registered up front sets `registers_on_demand = + False`, carves endpoints out of its pool in `allocate`, and revokes the pool in `shutdown`. +- **Batch mode**: a driver that registers regions together sets `registers_together = True` + on top of pool mode and implements `make_together_registration`, + `make_failing_together_registrations` and `make_busy_together_registration`. +- **Device memory**: override `allocate(nbytes, fill, memory_kind)` to allocate device memory + when the kind is DEVICE, and declare `device_memory=True`. `allocate` returns a pair: the + memory, an object implementing `Memory` (`address`, `nbytes`, `read`, `write`, + `canaries_intact`), and its registration handle when the driver registered it up front, or + `None` when the suite is to register it. The same checks then run over GPU memory; nothing + else changes. A sketch over a CUDA tensor, with the protocol's signatures: + + ```python + class DeviceBuffer: + def __init__(self, nbytes: int, fill: int) -> None: + self._tensor = torch.full((nbytes,), fill, dtype=torch.uint8, device="cuda") + + address = property(lambda self: self._tensor.data_ptr()) + nbytes = property(lambda self: self._tensor.numel()) + + def read(self, offset: int = 0, nbytes: int | None = None) -> bytes: + end = self.nbytes if nbytes is None else offset + nbytes + return bytes(self._tensor[offset:end].cpu().numpy()) + + def write(self, data: bytes, offset: int = 0) -> None: + source = torch.frombuffer(bytearray(data), dtype=torch.uint8) + self._tensor[offset : offset + len(data)].copy_(source) + + def canaries_intact(self) -> bool: + return True # a sketch: no guard bytes, so a write past the region goes unseen + + + class TestMyDriverConformance(DriverConformance): + def allocate(self, nbytes, fill, memory_kind): + if memory_kind is MemoryKind.DEVICE: + return DeviceBuffer(nbytes, fill), None # the suite registers it + return HostBuffer(nbytes, fill), None + ``` + +- **Proving the checks bite**: `testing/broken_drivers.py` pairs each check with a driver that + breaks exactly its rule; `test_memory_driver.py` runs every pair. A new check comes with a + new pair. + +## Assembling a backend + +```python +driver = open_driver(store_config) # whoever opens the driver closes it +binding = RegionBinding(regions, driver, hold) # hold: the lender's region hold +binding.register() # ConstructionRejected: nothing registered, the hold released +backend = BlobBackend(binding, wake, BlobConfig(key_prefix="trtllm-kv")) + +... # the lease holder fetches and publishes + +labels = backend.shutdown(timeout) # what is still accessible, if anything +if not binding.live_labels(): + binding.release_hold() # on the manager's thread + driver.close() # after the drain and the revocation +``` + +`register` raises `ConstructionRejected` only when nothing is registered and the hold is +already released. Any other raise means a registration may be live: keep the binding, call +`try_unregister_all()` again later, and `release_hold()` once `live_labels()` is empty. The +same holds when `BlobBackend(...)` itself raises, as when its progress thread cannot start: the +binding its owner registered is still live, and that owner revokes and releases as after any +failed build. + +Whoever opened the driver closes it, and owns it until then: after `shutdown` returned no +labels and `release_hold()` succeeded, the driver's `close()` is the last step. While anything +is still accessible, the owner keeps the driver, the binding and the hold alive; nothing frees +them on its own. A store that takes all of its memory in one step revokes it in one step too, +through `RegistersRegionsTogether`, and the same rule applies to the batch. + +`shutdown` returns the labels of what the backend still holds: `attempt:` for an attempt +whose access has not ended, `region:` for a region still registered. It returns early once +nothing within its budget could change that, and stops nothing: the progress thread keeps +retrying the drain and the revocation, and the hold stays with the manager until +`release_hold()` succeeds. A non-empty set means "keep everything, ask again later", never +"free it". + +## Two examples + +1. **Publish and fetch between two endpoints**: + `test_blob_example.py::test_two_endpoints_share_objects_through_the_store`. The writer's + source bytes are untouched by its publish; the reader's miss leaves its slot alone; the + store keeps its own copy, so the source may change once the publish quiesced. +2. **A lease holder's loop**: `FakeLeaseHolder` in `cpu_fixtures.py`, and the tests around it. + Decide once, from the outcome or from nothing at the deadline; return the lease on + `quiesced()` alone; release the region hold only when no region is registered. + +## Decisions worth knowing + +- **Batch registration**: `register_regions` registers a set of regions under one handle + and promises no atomicity. `DriverRejected` or `RegistrationError(handle=None)` certify + nothing is registered; `RegistrationError(handle=h)` says some of it may be and `h` revokes + whatever is; any other raise means some of it may be registered with nothing to revoke it + by, and the binding keeps the hold for good. +- **Who closes the driver**: the one who opened it, after `shutdown` returned no labels and the + hold went back. The backend never calls `close`; a cleanup that did not finish keeps the owner, + the driver and the memory, and is retried, never skipped. +- **`RegistrationError(handle=None)`** means the store certainly did not register the + region. A driver that cannot tell gives the handle. A failure it has no handle for is not a + `RegistrationError(None)`: the original exception escapes and the caller keeps the memory + protected. +- **Hints** reach the driver only on a fetch and only when it declares `hints`; a hint-free + fetch keeps the hint-free call shape. +- **Cancel** is asked of a driver that declares `cancel`, once per call of a stopped attempt + whose rows have not all concluded, by the progress thread; it is never access-end evidence. +- **No queue of its own**: every call goes to the driver as it is submitted, so the backend + bounds nothing and orders nothing. How many calls a store runs at once, in what order, and + when it refuses more (`DriverRejected`, which fails that call) is the driver's; how much to + submit at a time is the lease holder's admission policy. There is no backend-side cap on + queued work. +- **One call per group, cut twice**: `BlobBackend.driver_call_limit(region_index)` tells a + holder how many rows of a region one driver call carries (the driver's `max_rows_per_op`), + so it can bound what a single call that never proves its access ended would pin; `None` + means the driver declares no limit on the rows of one call. A group is also cut at a + repeated object key: a view may name one key at two addresses, while a `RowBatch` carries no + key twice. The rows keep their order, and the outcome's masks their positions. +- **A driver out of its protocol** (a raise from `put`/`get` that is not `DriverRejected`, no + handle returned, a raise from a query, a `poll` with the wrong row count) is fenced: the call + is held as unprovable unless its access end was recorded already, the calls of the same + submission after it never reach the driver, and every later submission is rejected. diff --git a/tensorrt_llm/_torch/disaggregation/blob/__init__.py b/tensorrt_llm/_torch/disaggregation/blob/__init__.py index efab37bb6b99..80de5d3d1a62 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/blob/__init__.py @@ -2,7 +2,17 @@ # 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``. +Read top down; a fetch or publish flows ``BlobBackend`` -> ``StoreDriver`` -> store SDK: + +- ``backend``: ``BlobBackend``, the public ``base.backend.Backend``. It checks each view, cuts it + into calls and hands each to a ``store.StoreDriver`` as one ``RowBatch``; the driver runs the + store's work on its own threads and tells the backend of changes through ``on_change``. + ``RegionBinding`` registers the lender's regions the driver can access, under the lender's hold, + one by one or, for a ``RegistersRegionsTogether`` driver, as one batch. +- ``store`` and ``keys``: the private driver boundary, which shares only ``MemoryKind`` with the + public contract, and the store key encoding. How an object is laid out in a store is its + driver's private format. + +No module imports torch, ``tensorrt_llm._torch.pyexecutor`` or a store SDK at import time. +``README.md`` next to this file says how to write a driver and how to assemble a backend. """ diff --git a/tensorrt_llm/_torch/disaggregation/blob/backend.py b/tensorrt_llm/_torch/disaggregation/blob/backend.py new file mode 100644 index 000000000000..b985d207a960 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/blob/backend.py @@ -0,0 +1,1070 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The blob backend: the public ``base.backend.Backend`` over a private store driver. + +The backend translates and keeps the books; the driver owns how its store's work runs: its queue, +the order in which gets and puts run, and its threads. A submission is checked and cut into calls, +one per layer group (split where the driver limits the rows of one call, and at a repeated object +key, which no call carries twice). A call is one driver operation, its ``OpHandle``; a driver may +split it further. Each call goes to the driver at once, on the caller's thread, as one ``put`` or +``get`` of a ``RowBatch``: the driver only records and queues it. The backend queues nothing, so +it has no order of its own to keep. + +Each call carries an ``on_change`` callback, which the driver calls once the operation's rows all +concluded, once its access ended and once it turned unprovable; it is the backend's only news of +the call. The callback queries the operation through the driver and records what the queries +show, on the thread the driver called it from. A record only ever adds evidence (results, an access +end, unprovability), so records may come in any order, on any thread, more than once. All +bookkeeping is under one lock; neither the driver nor ``wake`` is ever called under it. + +A call settles once its row results are recorded and its access end is recorded or can never be +proven ("unprovable"). An unprovable call holds its slots for good: its attempt never quiesces and +the regions stay registered. + +An attempt's outcome lands once every call has its results; it quiesces once every call's access +ended. A cancel, a fetch's first fault, a publish's whole-call failure or a shutdown stops it: when +the driver declares ``cancel``, the progress thread asks it to cancel the calls whose rows have not +all concluded. That, and the teardown after a shutdown (draining the driver at once, revoking the +registrations once every attempt quiesced), is all the progress thread does. + +A submission hands its calls over up to the first one the driver does not take, or until a +protocol violation is recorded; the calls after that never reach the driver. A driver that breaks +its protocol (a raise other than ``DriverRejected`` from ``put`` or ``get``, no handle, a raise +from a query, or a wrong result count) leaves its call held as unprovable, unless the call's access +end was recorded already, and every later submission is rejected. +""" + +from __future__ import annotations + +import functools +import itertools +import threading +import time +from collections.abc import Sequence +from dataclasses import dataclass +from typing import NamedTuple, NoReturn, Protocol, cast + +import numpy as np + +from tensorrt_llm.logger import logger + +from ..base.backend import ( + Capabilities, + Hint, + InvalidSubmission, + LayerGroupRows, + Outcome, + SlotRegion, + SubmissionRejected, + View, + WakeFn, +) +from ..base.memory import MemoryKind +from .keys import check_key_prefix, key_stem +from .store import ( + DriverRejected, + OnChange, + OpHandle, + RegionHandle, + RegionLayout, + RegionSpec, + RegistersRegionsTogether, + RegistrationError, + RowBatch, + RowResult, + RowStatus, + Span, + StoreDriver, + concluded, +) + +__all__ = [ + "BlobAttempt", + "BlobBackend", + "BlobConfig", + "ConstructionRejected", + "DriverCallLimit", + "RegionBinding", + "RegionHold", +] + +# After a shutdown, the progress thread retries what is left of the teardown (an attempt not +# quiesced yet, a drain that answered False, a revocation that failed) at an interval that starts +# at TEARDOWN_RETRY_MIN_S, doubles up to TEARDOWN_RETRY_MAX_S, and drops back to +# TEARDOWN_RETRY_MIN_S whenever the thread is poked. +TEARDOWN_RETRY_MIN_S = 0.001 +TEARDOWN_RETRY_MAX_S = 0.1 +SHUTDOWN_FAULT = "blob backend shut down" + +# Numbers each backend and binding: it names a backend's progress thread and keys each instance's +# warning_once calls, so one instance's warning never silences another's. +_instance_ids = itertools.count() + + +@dataclass(frozen=True) +class BlobConfig: + """The ``blob`` section of a store config. + + How a store's work is batched, scheduled and overlapped is its driver's own, so it is set with + the driver, for Mooncake in the ``mooncake`` section. + + Attributes: + key_prefix: The deployment's key namespace; matches ``[A-Za-z0-9_-]{1,64}``. + """ + + key_prefix: str = "trtllm-kv" + + def __post_init__(self) -> None: + if not isinstance(self.key_prefix, str): + raise ValueError(f"key_prefix must be a str, got {self.key_prefix!r}") + check_key_prefix(self.key_prefix) + + +class ConstructionRejected(RuntimeError): + """``RegionBinding.register`` failed and left nothing registered, held or accessible: no cleanup + is needed, and a backend factory maps it to the runtime's certified rejection. + + Any other exception from ``register`` may leave a registration live; the binding then keeps the + hold, and ``release_hold`` refuses until every registration is revoked. + """ + + +class RegionHold(Protocol): + """What this package uses of ``pyexecutor.kv_cache.sharing.RegionHold``, restated because blob + modules may not import ``pyexecutor``.""" + + def release(self) -> None: ... + + +class RegionBinding: + """The lender's region hold and the registrations made under it. + + Registers, once each and on the builder thread, the regions the driver can access; revokes the + registrations only when nothing can access them, on the backend's progress thread (or on the + builder thread when ``register`` rolls back a failed registration); releases the hold only once + no registration may be live: through ``release_hold`` on the manager's thread, or inside + ``register`` on the builder thread when it raises ``ConstructionRejected``, after which the + caller must not call ``release_hold``. The driver is never called under the binding's lock, so + reading the live registrations never waits on a revocation. + + A driver implementing ``RegistersRegionsTogether`` takes every region in one + ``register_regions`` call and gives them back in one ``unregister_regions``; every region has + an owner here before the store is asked, so a failure of any shape leaves them accounted for. + + Args: + regions: The lender's regions, in its order; region ``i`` is labelled ``region:``. + driver: The store driver that registers them. + hold: The lender's hold, taken before registering. + """ + + def __init__( + self, regions: Sequence[SlotRegion], driver: StoreDriver, hold: RegionHold + ) -> None: + self.regions = tuple(regions) + self.driver = driver + device_memory = driver.capabilities.device_memory + # The regions register() registers, by index: those with slots in memory the driver can + # access. A row in any other region is one the backend does not recognize. + self.served_region_indices = frozenset( + index + for index, region in enumerate(self.regions) + if region.slot_count > 0 and (region.memory_kind is MemoryKind.HOST or device_memory) + ) + self._hold = hold + # The driver as a RegistersRegionsTogether, taking the regions as one batch; else None. + self._together = driver if isinstance(driver, RegistersRegionsTogether) else None + self._name = f"region-binding-{next(_instance_ids)}" + self._lock = threading.Lock() + # Region index -> its handle, or None for a registration that may exist without one. + self._live_registrations: dict[int, RegionHandle | None] = {} + + def register(self) -> None: + """Registers the regions in ``served_region_indices``. + + Raises: + ConstructionRejected: Nothing is registered and the hold is released: there is no region + to register, a region's layout is refused, or a registration failed with nothing + escaping and the earlier ones were revoked. + RuntimeError: A registration may be live; the hold is kept. + """ + if not self.served_region_indices: # a backend that could serve no row + self._hold.release() + raise ConstructionRejected("no region has slots in memory the store driver can access") + layouts = {} + for index in sorted(self.served_region_indices): + region = self.regions[index] + try: # every layout is checked before any region is registered + layouts[index] = RegionLayout( + region.slot_bytes, region.slot_count, region.layout_label, region.memory_kind + ) + except ValueError as error: + self._hold.release() + raise ConstructionRejected(f"region:{index}: {error}") from error + if self._together is not None: + self._register_together(self._together, layouts) + return + for index, layout in layouts.items(): + region = self.regions[index] + try: + handle = self.driver.register(region.address, region.nbytes, layout) + except Exception as error: + self._roll_back_if_possible_and_raise(index, error) + else: + self._track_live_registration(index, handle) + + def try_unregister_all(self) -> bool: + """Revokes every live registration it can. True when none is left; a revocation that failed + stays live for the next call.""" + if self._together is not None: + return self._unregister_together(self._together) + with self._lock: + revocable = [ + (i, handle) for i, handle in self._live_registrations.items() if handle is not None + ] + for index, handle in revocable: + try: + self.driver.unregister(handle) + except Exception as error: + logger.warning_once( + f"blob backend: revoking region:{index} failed, will retry: {error!r}", + key=(self._name, "unregister", index), + ) + continue + with self._lock: + del self._live_registrations[index] + with self._lock: + return not self._live_registrations + + def release_hold(self) -> None: + """Releases the lender's hold, on the manager's thread. + + Raises: + RuntimeError: A registration may still be live; the hold is kept. + """ + live = self.live_labels() + if live: + raise RuntimeError(f"the region hold is kept: {sorted(live)} may still be registered") + self._hold.release() + + def live_labels(self) -> frozenset[str]: + """``region:`` for every region that may still be registered.""" + with self._lock: + return frozenset(f"region:{index}" for index in self._live_registrations) + + def _roll_back_if_possible_and_raise(self, index: int, error: Exception) -> NoReturn: + """Raises for region ``index``'s failed registration. When nothing escaped, unregisters + every earlier region; once none is left, releases the lender's hold and raises + ``ConstructionRejected``. Otherwise (the failed registration may be live, or an earlier + region stays registered) keeps the hold and raises ``RuntimeError``.""" + handle = error.handle if isinstance(error, RegistrationError) else None + nothing_escaped = isinstance(error, DriverRejected) or ( + isinstance(error, RegistrationError) and handle is None + ) + if not nothing_escaped: + self._track_live_registration(index, handle) + raise RuntimeError(f"region:{index} may be registered after {error!r}") from error + if not self.try_unregister_all(): + raise RuntimeError( + f"registering region:{index} failed and an earlier region stays registered: {error}" + ) from error + # Nothing is registered or accessible, and construction precedes the executor loop, so the + # builder thread may still give the hold back. + self._hold.release() + raise ConstructionRejected(f"registering region:{index} failed: {error}") from error + + def _track_live_registration(self, index: int, handle: RegionHandle | None) -> None: + with self._lock: + self._live_registrations[index] = handle + + def _register_together( + self, driver: RegistersRegionsTogether, layouts: dict[int, RegionLayout] + ) -> None: + try: + specs = [ + RegionSpec(self.regions[index].address, self.regions[index].nbytes, layout) + for index, layout in layouts.items() + ] + except ValueError as error: # a region the boundary cannot describe; the store not asked + self._hold.release() + raise ConstructionRejected(f"bad region geometry: {error}") from error + with self._lock: # every region has an owner before the store is asked + self._live_registrations = dict.fromkeys(layouts) + try: + handle = driver.register_regions(specs) + except Exception as error: + self._raise_for_failed_batch(error) + self._keep_all(handle) + + def _unregister_together(self, driver: RegistersRegionsTogether) -> bool: + with self._lock: + handles: list[RegionHandle] = [] + for handle in self._live_registrations.values(): + if handle is not None and not any(handle is seen for seen in handles): + handles.append(handle) + for handle in handles: # the one batch handle every region shares + try: + driver.unregister_regions(handle) + except Exception as error: + logger.warning_once( + f"blob backend: revoking the regions together failed, will retry: {error!r}", + key=(self._name, "unregister_regions"), + ) + continue + with self._lock: + self._live_registrations = { + i: h for i, h in self._live_registrations.items() if h is not handle + } + with self._lock: + return not self._live_registrations + + def _raise_for_failed_batch(self, error: Exception) -> NoReturn: + handle = error.handle if isinstance(error, RegistrationError) else None + nothing_escaped = isinstance(error, DriverRejected) or ( + isinstance(error, RegistrationError) and handle is None + ) + if nothing_escaped: + with self._lock: + self._live_registrations.clear() + self._hold.release() + raise ConstructionRejected( + f"registering the regions together failed: {error}" + ) from error + if handle is None: # some may be registered, with nothing to revoke them by + raise RuntimeError(f"the regions may be registered after {error!r}") from error + self._keep_all(handle) + if not self.try_unregister_all(): + raise RuntimeError( + f"registering the regions together failed and some may stay registered: {error}" + ) from error + self._hold.release() + raise ConstructionRejected(f"registering the regions together failed: {error}") from error + + def _keep_all(self, handle: RegionHandle) -> None: + with self._lock: + for index in self._live_registrations: + self._live_registrations[index] = handle + + +class DriverCallLimit(NamedTuple): + """The most rows of a region, and their bytes, that one driver ``put`` or ``get`` carries.""" + + rows: int + nbytes: int + + +class _GroupRows(NamedTuple): + """One group's arrays, taken from the view at submission.""" + + object_keys: np.ndarray # uint8 [rows, width] + addresses: np.ndarray # int64 [rows] + region_index: int + + +@dataclass(eq=False) +class _Call: + """Rows ``[start, stop)`` of one group of an attempt, handed to the driver as one operation: one + ``put`` or ``get``, one ``OpHandle``.""" + + attempt: BlobAttempt + group_position: int # in ``View.groups``, not the group's ``layer_group_index`` + start: int + stop: int + # The fields below change only under the backend's lock. + op: OpHandle | None = None # the driver's handle, once the whole submission was handed over + results_recorded: bool = False + access_end_recorded: bool = False + unprovable: bool = False # the driver can never prove its access ended + cancel_forwarded: bool = False # the driver was asked to cancel it + + def __len__(self) -> int: + return self.stop - self.start + + +class _CallResults(NamedTuple): + """A concluded call's results as an attempt records them: whether each row is READY, and the + first row that faulted (its index in the call) with its reason.""" + + ready: list[bool] + first_fault: tuple[int, str] | None + + @classmethod + def none_delivered(cls, rows: int, fault: str | None = None) -> _CallResults: + """Every row undelivered: each faults with ``fault`` when given (the call failed as a + whole), otherwise none does (CANCELLED).""" + return cls([False] * rows, None if fault is None else (0, fault)) + + @classmethod + def of(cls, results: Sequence[RowResult]) -> _CallResults: + ready = [result.status is RowStatus.READY for result in results] + if all(ready): + return cls(ready, None) + faults = ((row, r.fault) for row, r in enumerate(results) if r.fault is not None) + return cls(ready, next(faults, None)) + + +class _Recorded(NamedTuple): + """What a record under the backend's lock leaves to do once the lock is released.""" + + wake_due: bool # the attempt's outcome landed or it quiesced + newly_unprovable: bool # the call turned unprovable with this record + + +class BlobAttempt: + """The handle of one submission: a ``base.backend.Attempt``. + + Its state changes only under the backend's lock; ``poll`` and ``quiesced`` read one reference + without it, so they never block. + + Two counters, both set to the number of calls at submission; each call decrements each at most + once: + + - ``_calls_awaiting_results``: calls whose row results are not recorded yet. + - ``_calls_awaiting_access_end``: calls whose access end is not recorded yet. An unprovable + call never leaves it. + + The outcome lands once ``_calls_awaiting_results`` is 0; the attempt quiesces once + ``_calls_awaiting_access_end`` is 0. + """ + + def __init__( + self, + backend: BlobBackend, + label: str, + is_fetch: bool, + groups: list[_GroupRows], + call_ranges: list[tuple[int, int, int]], + hint: Hint = None, + ) -> None: + self.label = label + """The attempt's label among ``BlobBackend.shutdown``'s, ``attempt:``.""" + self.is_fetch = is_fetch + self._backend = backend + self._hint = hint # a fetch's source, for every driver get; immutable, so kept as given + self._masks = [np.zeros(len(rows.addresses), dtype=bool) for rows in groups] + self._calls = [_Call(self, *call_range) for call_range in call_ranges] + self._calls_awaiting_results = self._calls_awaiting_access_end = len(self._calls) + # The first row fault; _record_results decides whether it stops the attempt. + self._first_fault: str | None = None + # Set by a cancel, a stopping fault or a shutdown; when the driver declares cancel, the + # progress thread then asks it to cancel the calls not concluded. + self._stopped = False + self._outcome: Outcome | None = None + self._quiesced = False + + def poll(self) -> Outcome | None: + return self._outcome + + def quiesced(self) -> bool: + return self._quiesced + + def cancel(self) -> None: + """Stops the attempt unless its outcome landed. When the driver declares ``cancel``, the + progress thread asks it to cancel the calls whose rows have not all concluded; rows already + delivered stay delivered.""" + self._backend._cancel(self) + + @property + def _has_unprovable_call(self) -> bool: + """A call can never prove its access ended, so the attempt never quiesces.""" + return any(call.unprovable for call in self._calls) + + def _record_results(self, call: _Call, results: _CallResults, whole_call_failed: bool) -> bool: + """Writes ``call``'s results into the masks and keeps the first fault. True when that fault + stops the attempt: when the driver declares ``cancel``, every call not concluded yet is + cancelled, earlier ones included, and rows already delivered stay delivered. Any fault stops + a fetch: the request it serves has a gap to fill some other way. A publish's rows stand + alone, so only a call that failed as a whole stops it.""" + self._masks[call.group_position][call.start : call.stop] = results.ready + if results.first_fault is None: + return False + row, fault = results.first_fault + where = f"group {call.group_position} row {call.start + row}" # indexes the outcome's masks + self._first_fault = self._first_fault or f"{where}: {fault}" + return self.is_fetch or whole_call_failed + + def _conclude_or_quiesce(self) -> bool: + """Lands the outcome and turns quiesced once their conditions hold; True when either did.""" + changed = False + # The outcome waits for every call's results even after a fault: calls conclude in any + # order, so the first fault seen can come from a call issued after one that is still + # delivering rows a prefix reuse needs. + if self._outcome is None and self._calls_awaiting_results == 0: + for mask in self._masks: # final, so the outcome keeps them without a copy + mask.flags.writeable = False + self._outcome = Outcome(tuple(self._masks), self._first_fault) + changed = True + if not self._quiesced and self._calls_awaiting_access_end == 0: + self._quiesced = True + changed = True + return changed + + def _force_outcome(self, reason: str) -> None: + first = self._first_fault + fault = reason if first is None else f"{reason}; first fault: {first}" + # Still writable, so ``Outcome`` copies them: calls still open go on writing their results + # into these masks. + self._outcome = Outcome(tuple(self._masks), fault) + + +class BlobBackend: + """Moves a view's rows between the lender's registered regions and a byte store: a + ``base.backend.Backend``. A process may hold several, one per lender, each with its own driver. + + Args: + binding: The lender's regions, already registered with ``binding.driver``. The backend + revokes the registrations once nothing can access them after shutdown. + wake: Called outside every lock whenever an attempt's outcome lands or it quiesces, once + the change shows through ``poll`` or ``quiesced``. + config: The store config's ``blob`` section. + """ + + def __init__(self, binding: RegionBinding, wake: WakeFn, config: BlobConfig) -> None: + # Fixed once __init__ returns. + self._binding = binding + self._driver = binding.driver + self._regions = binding.regions + # Each region's whole extent, which a RowBatch names as holding all its rows. + self._region_extents = { + index: Span(region.address, region.nbytes) + for index, region in enumerate(binding.regions) + if index in binding.served_region_indices + } + self._wake = wake + self._key_stem = key_stem(config.key_prefix) + self._name = name = f"blob-backend-{next(_instance_ids)}" + driver_capabilities = self._driver.capabilities + self._max_rows_per_op = driver_capabilities.max_rows_per_op + self._driver_reads_hints = driver_capabilities.hints + self._driver_cancels = driver_capabilities.cancel + # Every call goes to the driver as it is submitted, so the backend can cancel only by + # asking the driver: it declares cancel exactly when the driver does. + self._capabilities = Capabilities( + fetch=driver_capabilities.get, + publish=driver_capabilities.put, + cancel=driver_capabilities.cancel, + ) + self._attempt_ids = itertools.count() # its next() is atomic: submissions take it unlocked + self._lock = threading.Lock() + # The progress thread waits on this between rounds; _poke_progress notifies it. + self._progress_due = threading.Condition(self._lock) + # ``shutdown`` waits on this; notified whenever its answer may have changed: an attempt + # concluded or quiesced, or the progress thread finished the teardown. + self._shutdown_recheck = threading.Condition(self._lock) + + # Guarded by _lock. + # Set by _poke_progress, cleared as the next round starts: run a round now. + self._progress_poked = False + self._attempts: set[BlobAttempt] = set() # until concluded and quiesced + # Set once and never cleared. The progress loop also reads it without the lock: a round + # that misses it runs again, since setting it pokes the progress thread. + self._shutdown_requested = False + self._teardown_done = False # drained and every registration revoked + self._protocol_violation: str | None = None # how the driver first broke its protocol + + # Used by the progress thread only. + self._drained = False # the driver's drain returned True; it is not drained again + # Used by ``shutdown`` only: the labels the last shutdown logged, which it logs again only + # once they changed. + self._last_logged_labels: frozenset[str] | None = None + + self._progress_thread = threading.Thread( + target=self._progress_loop, name=f"{name}-progress", daemon=True + ) + self._progress_thread.start() + + @property + def capabilities(self) -> Capabilities: + return self._capabilities + + def fetch(self, view: View, source: Hint) -> BlobAttempt: + """``source`` goes unchanged to every driver ``get`` of the fetch when the driver declares + ``hints``; any other driver takes only None.""" + return self._submit(view, source, is_fetch=True) + + def publish(self, view: View, destination: Hint) -> BlobAttempt: + """``destination`` must be None: a driver ``put`` takes no hint.""" + return self._submit(view, destination, is_fetch=False) + + def shutdown(self, timeout: float) -> frozenset[str]: + """Forces outcomes, then waits up to ``timeout`` for the progress thread to finish cleanly, + or only until the answer can no longer change: every attempt not yet quiesced holds a call + whose access can never be proven to end. + + Never calls the driver: the progress thread forwards cancels, drains the driver and, once + every attempt quiesced, revokes the registrations, since revocation can be slow (RDMA + deregisters on every NIC) and a shutdown with no budget left must still return at once. + """ + # Bounded so that the waits below accept it: an infinite budget must not raise. + deadline = time.monotonic() + min(max(0.0, timeout), threading.TIMEOUT_MAX) + if self._refuse_new_work_and_force_outcomes(): + self._call_wake() + with self._lock: + self._shutdown_recheck.wait_for( + self._shutdown_answer_is_final, max(0.0, deadline - time.monotonic()) + ) + teardown_done = self._teardown_done + if teardown_done: # the progress thread is returning + self._progress_thread.join(max(0.0, deadline - time.monotonic())) + labels = self._still_accessible_labels() + self._log_still_accessible(labels) + return labels + + def driver_call_limit(self, region_index: int) -> DriverCallLimit | None: + """The most rows of region ``region_index`` the driver takes in one ``put`` or ``get`` (its + ``max_rows_per_op``), and their bytes; None when the driver declares no limit on the rows + of one call. A group is still cut at a repeated object key, which no call carries twice. + + An attempt quiesces only once every call proved its access ended, so one call that never + proves it keeps the slots of its whole submission; a holder bounds that by submitting at + most this many rows of the region at a time. + + Raises: + ValueError: The backend did not register the region. + """ + if region_index not in self._binding.served_region_indices: + raise ValueError(f"region {region_index} is not registered") + rows = self._max_rows_per_op + if rows is None: + return None + return DriverCallLimit(rows=rows, nbytes=rows * self._regions[region_index].slot_bytes) + + def stats(self) -> dict[str, int]: + """A snapshot of the work in flight, for logs and tests.""" + with self._lock: + calls = [call for attempt in self._attempts for call in attempt._calls] + return { + "attempts_active": len(self._attempts), + "calls_awaiting_results": sum(not call.results_recorded for call in calls), + "calls_awaiting_access_end": sum(not call.access_end_recorded for call in calls), + "calls_unprovable": sum(call.unprovable for call in calls), + } + + # Submission, on the executor loop's thread: validates, then hands each call to the driver, + # which only records and queues it. + + def _submit(self, view: View, hint: Hint, is_fetch: bool) -> BlobAttempt: + if hint is not None and not is_fetch: + raise InvalidSubmission(f"a publish takes only a None destination, got {hint!r}") + if hint is not None and not self._driver_reads_hints: + raise InvalidSubmission(f"the store driver reads no hints, got {hint!r}") + if not (self._capabilities.fetch if is_fetch else self._capabilities.publish): + raise InvalidSubmission(f"the store driver cannot {'get' if is_fetch else 'put'}") + groups = [self._checked_group_rows(group) for group in view.groups] + if is_fetch: + _refuse_shared_destinations(groups) + label = f"attempt:{next(self._attempt_ids)}" + attempt = BlobAttempt(self, label, is_fetch, groups, self._call_ranges(groups), hint) + batches = [self._row_batch(groups, call) for call in attempt._calls] + with self._lock: + if self._shutdown_requested: + raise SubmissionRejected("the blob backend is shut down") + if self._protocol_violation is not None: + raise SubmissionRejected( + f"the store driver broke its protocol: {self._protocol_violation}" + ) + self._attempts.add(attempt) + wake_due = self._conclude_or_quiesce_and_retire(attempt) # an empty view ends here + if wake_due: + self._call_wake() + self._hand_over(attempt, batches) + return attempt + + def _checked_group_rows(self, group: LayerGroupRows) -> _GroupRows: + where = f"layer group {group.layer_group_index}" + if group.object_keys is None or group.addresses is None or group.region_index is None: + raise InvalidSubmission(f"{where}: rows need object keys, addresses and a region") + keys = np.asarray(group.object_keys) + addresses = np.asarray(group.addresses) + rows = len(group) + if keys.dtype != np.uint8 or keys.ndim != 2 or keys.shape[0] != rows or keys.shape[1] == 0: + raise InvalidSubmission( + f"{where}: object keys must be uint8 [{rows}, width > 0], got {keys.dtype} " + f"{keys.shape}" + ) + keys = np.ascontiguousarray(keys) # one key twice is allowed: _call_bounds cuts there + if addresses.dtype != np.int64 or addresses.shape != (rows,): + raise InvalidSubmission( + f"{where}: addresses must be int64 [{rows}], got {addresses.dtype} " + f"{addresses.shape}" + ) + # Refused, never skipped: the driver must not be handed memory it never registered. + if group.region_index not in self._binding.served_region_indices: + raise InvalidSubmission(f"{where}: region {group.region_index} is not registered") + region = self._regions[group.region_index] + # Known limit: numpy lets go of the GIL for arrays of more than about 500 elements, so a + # group that large may wait here behind a busy thread. + slot, misalignment = np.divmod(addresses - region.address, region.slot_bytes) + if np.any((slot < 0) | (slot >= region.slot_count) | (misalignment != 0)): + raise InvalidSubmission(f"{where}: an address is not a slot of its region") + return _GroupRows(keys, addresses, group.region_index) + + def _call_ranges(self, groups: list[_GroupRows]) -> list[tuple[int, int, int]]: + """``(group position, start, stop)`` of each call: one per group, split where the driver + limits the rows of one call and at a repeated object key.""" + ranges = [] + for position, group_rows in enumerate(groups): + rows_per_call = self._max_rows_per_op or max(len(group_rows.addresses), 1) + ranges += [ + (position, start, stop) + for start, stop in _call_bounds(group_rows.object_keys, rows_per_call) + ] + return ranges + + def _row_batch(self, groups: list[_GroupRows], call: _Call) -> RowBatch: + group_rows = groups[call.group_position] + rows = slice(call.start, call.stop) + region_index = group_rows.region_index + return RowBatch.of_slots( + self._key_stem, + group_rows.object_keys[rows], + group_rows.addresses[rows], + self._regions[region_index].slot_bytes, + self._region_extents[region_index], + ) + + def _hand_over(self, attempt: BlobAttempt, batches: list[RowBatch]) -> None: + """Hands ``attempt``'s calls to the driver in order, up to the first one the driver does + not take, or the first after a protocol violation was recorded (an earlier call's + ``on_change`` may record one meanwhile); that call fails as a whole, and the calls after it + never reach the driver and conclude CANCELLED.""" + taken: list[tuple[_Call, OpHandle]] = [] + for call, rows in zip(attempt._calls, batches): + # Read without the lock: a violation recorded a moment later lets one more call + # through, which is then held like any call already in the driver. + violation = self._protocol_violation + if violation is not None: + reason = f"not handed over: the store driver broke its protocol ({violation})" + self._fail_call(call, reason, this_call_broke_protocol=False) + break + op = self._hand_over_call(call, rows) + if op is None: + break + taken.append((call, op)) + for call in attempt._calls[len(taken) + 1 :]: # past the one the driver did not take + self._record(call, _CallResults.none_delivered(len(call)), access_ended=True) + with self._lock: + for call, op in taken: + call.op = op + # The progress thread cancels only calls whose handle it knows, so a stop that came + # while the calls were being handed over is forwarded from here. + if attempt._stopped and self._driver_cancels: + self._poke_progress() + + def _hand_over_call(self, call: _Call, rows: RowBatch) -> OpHandle | None: + """The driver's handle for ``call``, or None when the driver did not take it; the failure + is then recorded.""" + on_change: OnChange = functools.partial(self._query_and_record, call) + attempt = call.attempt + try: + if not attempt.is_fetch: + op = self._driver.put(rows, on_change=on_change) + elif attempt._hint is None: + op = self._driver.get(rows, on_change=on_change) + else: + op = self._driver.get(rows, attempt._hint, on_change=on_change) + except DriverRejected as error: # certifies that nothing escaped + self._fail_call(call, f"refused: {error}", this_call_broke_protocol=False) + return None + except Exception as error: # outside the driver protocol: the rows may be accessed for good + self._fail_call( + call, f"the store driver raised {error!r}", this_call_broke_protocol=True + ) + return None + if op is None: # outside the driver protocol too: nothing to query or cancel + self._fail_call( + call, "the store driver returned no handle", this_call_broke_protocol=True + ) + return op + + # Recording, on the thread that calls ``on_change``, or on the submitting one for a call the + # driver did not take. + + def _query_and_record(self, call: _Call, op: OpHandle) -> None: + """``call``'s ``on_change``: queries its operation through the driver and records what the + queries show. A record can stop the attempt, land its outcome and quiescence, and latch a + protocol violation.""" + try: + results = self._driver.poll(op) + access_ended = self._driver.access_ended(op) + unprovable = not access_ended and self._driver.access_unprovable(op) + count, all_concluded = len(results), concluded(results) + except Exception as error: # the queries never raise for a handle the driver issued + self._fail_call( + call, f"a store driver query raised {error!r}", this_call_broke_protocol=True + ) + return + if count != len(call): + reason = f"the store driver gave {count} results for {len(call)} rows" + self._fail_call(call, reason, this_call_broke_protocol=True) + return + # Read row by row here, outside the lock the submitting thread takes. + call_results = _CallResults.of(results) if all_concluded else None + self._record(call, call_results, access_ended, unprovable) + + def _fail_call(self, call: _Call, reason: str, this_call_broke_protocol: bool) -> None: + """Records ``call`` as failed as a whole: every row faults with ``reason``. + + A call the driver refused, or never saw, touched nothing, so its access ended. A call on + which the driver broke its protocol is held as unprovable (unless its access end was + recorded already): nothing the driver says about it later is taken as evidence. Calls + already in the driver keep being recorded from its answers, and every later submission is + rejected.""" + with self._lock: + first_violation = this_call_broke_protocol and self._protocol_violation is None + if first_violation: + self._protocol_violation = reason + recorded = self._record_locked( + call, + _CallResults.none_delivered(len(call), reason), + access_ended=not this_call_broke_protocol, + unprovable=this_call_broke_protocol, + whole_call_failed=True, + ) + if first_violation: + logger.warning( + f"blob backend: the store driver broke its protocol ({reason}); every later " + "submission is rejected" + ) + self._after_record(call.attempt, recorded) + + def _record( + self, + call: _Call, + results: _CallResults | None, + access_ended: bool, + unprovable: bool = False, + ) -> None: + with self._lock: + recorded = self._record_locked(call, results, access_ended, unprovable) + self._after_record(call.attempt, recorded) + + def _after_record(self, attempt: BlobAttempt, recorded: _Recorded) -> None: + """What a record leaves to do once the lock is released: a log line for a call that turned + unprovable, and the holder's ``wake``. Neither runs under the lock, so a slow log sink or + ``wake`` never holds up the submitting thread.""" + if recorded.newly_unprovable: + logger.warning_once( + f"blob backend: a store call of {attempt.label} can never prove its access " + "ended; its slots stay held", + key=(self._name, "unprovable"), + ) + if recorded.wake_due: + self._call_wake() + + # Bookkeeping, under the lock. + + def _record_locked( + self, + call: _Call, + results: _CallResults | None, + access_ended: bool, + unprovable: bool, + whole_call_failed: bool = False, + ) -> _Recorded: + """Records ``call``'s results (None while not all concluded), its access end and whether + that end can never be proven. Only adds evidence, so a stale or repeated record changes + nothing; a recorded access end stands against a later claim that it can never be proven, + and an unprovable call stays held whatever the driver says of it later. The caller passes + what it returns to ``_after_record`` once it released the lock.""" + attempt = call.attempt + if results is not None and not call.results_recorded: + call.results_recorded = True + attempt._calls_awaiting_results -= 1 + fault_stops_attempt = attempt._record_results(call, results, whole_call_failed) + if fault_stops_attempt: + self._stop(attempt) + if access_ended and not call.access_end_recorded and not call.unprovable: + call.access_end_recorded = True + attempt._calls_awaiting_access_end -= 1 + newly_unprovable = unprovable and not call.unprovable and not call.access_end_recorded + if newly_unprovable: + call.unprovable = True + return _Recorded(self._conclude_or_quiesce_and_retire(attempt), newly_unprovable) + + def _conclude_or_quiesce_and_retire(self, attempt: BlobAttempt) -> bool: + """Under the lock: lands ``attempt``'s outcome and quiescence once due, and drops it from + ``_attempts`` once it has both. True when the caller is to call ``wake``.""" + wake_due = attempt._conclude_or_quiesce() + if attempt._outcome is not None and attempt._quiesced: + self._attempts.discard(attempt) + if self._shutdown_requested: # the teardown may be one round away + self._poke_progress() + # Every attempt change a shutdown waits for passes through here; the progress thread + # notifies the end of the teardown itself. + self._shutdown_recheck.notify_all() + return wake_due + + def _cancel(self, attempt: BlobAttempt) -> None: + """Stops ``attempt`` unless its outcome already landed. Stopping changes no counter, so it + lands neither the outcome nor quiescence itself.""" + with self._lock: + if attempt._outcome is None: + self._stop(attempt) + + def _stop(self, attempt: BlobAttempt) -> None: + """Under the lock: none of ``attempt``'s calls is worth finishing any more. When the driver + declares ``cancel``, the progress thread asks it to cancel those still running.""" + if not attempt._stopped: + attempt._stopped = True + if self._driver_cancels: + self._poke_progress() + + def _poke_progress(self) -> None: + self._progress_poked = True + self._progress_due.notify() + + # The progress thread: forwards cancels and, after a shutdown, tears down. + + def _progress_loop(self) -> None: + retry_s = TEARDOWN_RETRY_MIN_S + while True: + try: + self._forward_cancels() + if self._shutdown_requested and self._try_teardown(): + with self._lock: + self._teardown_done = True + self._shutdown_recheck.notify_all() + return + except Exception as error: # the thread must live on, or the labels never clear + logger.warning_once( + f"blob backend: a progress round raised {error!r}", + key=(self._name, "progress", type(error)), + ) + retry_s = self._wait_for_next_round(retry_s) + + def _forward_cancels(self) -> None: + """Asks the driver, once per call, to cancel the calls of stopped attempts whose rows have + not all concluded.""" + if not self._driver_cancels: + return + with self._lock: + due = [ + call + for attempt in self._attempts + if attempt._stopped + for call in attempt._calls + if call.op is not None and not call.results_recorded and not call.cancel_forwarded + ] + for call in due: + call.cancel_forwarded = True + for call in due: + try: + self._driver.cancel(call.op) + except Exception as error: # advisory: the call runs on + logger.warning_once( + f"blob backend: cancelling a store call raised {error!r}", + key=(self._name, "cancel", type(error)), + ) + + def _wait_for_next_round(self, retry_s: float) -> float: + """Waits until poked or, while a teardown may still finish, for at most ``retry_s``; returns + the retry interval for the wait after it.""" + with self._lock: + if not self._progress_poked: + tearing_down = self._shutdown_requested and not self._shutdown_answer_is_final() + self._progress_due.wait(retry_s if tearing_down else None) + poked, self._progress_poked = self._progress_poked, False + return TEARDOWN_RETRY_MIN_S if poked else min(2 * retry_s, TEARDOWN_RETRY_MAX_S) + + def _try_teardown(self) -> bool: + """True once the driver is drained and the registrations revoked; a step that is not done + yet is retried next round. The driver is drained from the first round on, as nothing is + submitted any more: queued work is cancelled and the driver's idle threads exit, even when + an attempt never quiesces. The registrations are revoked only once every attempt quiesced. + The driver is never closed here: closing a store client may block for long, and it is + the step of the owner that opened it, once shutdown returned no labels and the hold went + back.""" + if not self._drained: + self._drained = self._driver.drain(0.0) + with self._lock: + if any(not attempt._quiesced for attempt in self._attempts): + return False + return self._drained and self._binding.try_unregister_all() + + # Shutdown, on the shutdown thread. + + def _refuse_new_work_and_force_outcomes(self) -> bool: + """Refuses every later submission, forces the shutdown outcome on each attempt not + concluded and stops it (its calls are cancelled when the driver declares cancel), and pokes + the progress thread to start the teardown. True when an outcome was forced, so the caller + calls ``wake``.""" + with self._lock: + self._shutdown_requested = True + forced = [attempt for attempt in self._attempts if attempt._outcome is None] + for attempt in forced: + self._stop(attempt) + attempt._force_outcome(SHUTDOWN_FAULT) + self._conclude_or_quiesce_and_retire(attempt) + self._poke_progress() # the progress thread forwards the cancels and tears down + return bool(forced) + + def _shutdown_answer_is_final(self) -> bool: + """Under the lock, once closing: the progress thread finished, or the labels are final. An + attempt with a call whose access can never be proven to end never quiesces, and the + registrations are revoked only once every attempt quiesced, so when every attempt not yet + quiesced has such a call, neither its label nor a region's can go any more.""" + if self._teardown_done: + return True + unquiesced = [attempt for attempt in self._attempts if not attempt._quiesced] + return bool(unquiesced) and all(attempt._has_unprovable_call for attempt in unquiesced) + + def _still_accessible_labels(self) -> frozenset[str]: + with self._lock: + attempts = {attempt.label for attempt in self._attempts if not attempt._quiesced} + return frozenset(attempts) | self._binding.live_labels() + + def _log_still_accessible(self, labels: frozenset[str]) -> None: + if labels == self._last_logged_labels: + return + self._last_logged_labels = labels + if labels: + logger.warning(f"blob backend: shutting down, still accessible: {sorted(labels)}") + else: + logger.info("blob backend: shut down; nothing is accessible any more") + + def _call_wake(self) -> None: + try: + self._wake() + except Exception as error: # a raising wake must not kill a backend thread + logger.warning_once( + f"blob backend: wake raised {error!r}", key=(self._name, "wake", type(error)) + ) + + +def _call_bounds(keys: np.ndarray, step: int) -> list[tuple[int, int]]: + """Consecutive ``[start, stop)`` runs of at most ``step`` rows, none naming a key twice: a view + may name one key at two addresses, while a ``RowBatch`` promises the driver no key twice (a get + of one key into two slots would answer READY for both and write one). Keys are compared as + bytes in a Python set, not with ``np.unique``, whose sort lets go of the GIL and leaves this + thread waiting a switch interval behind any busy thread. Distinct keys, the usual case, cut + into whole steps.""" + count = len(keys) + if count == 0: + return [] + rows = cast( + list[bytes], + np.ascontiguousarray(keys).view(np.dtype((np.void, keys.shape[1]))).ravel().tolist(), + ) + if len(set(rows)) == count: + return [(start, min(start + step, count)) for start in range(0, count, step)] + bounds: list[tuple[int, int]] = [] + seen: set[bytes] = set() + start = 0 + for index, key in enumerate(rows): + if index - start == step or key in seen: + bounds.append((start, index)) + start, seen = index, set() + seen.add(key) + bounds.append((start, count)) + return bounds + + +def _refuse_shared_destinations(groups: list[_GroupRows]) -> None: + """Two fetched rows in one slot would let one report READY over the other's bytes. A publish + only reads its slots, so only a fetch is checked. A Python set, not np.unique, whose sort lets + go of the GIL.""" + seen: set[int] = set() + for rows in groups: + seen.update(rows.addresses.tolist()) + if len(seen) != sum(len(rows.addresses) for rows in groups): + raise InvalidSubmission("two rows of a fetch write the same slot") diff --git a/tensorrt_llm/_torch/disaggregation/blob/keys.py b/tensorrt_llm/_torch/disaggregation/blob/keys.py index a27af1c7d481..82a6f9dfa546 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/keys.py +++ b/tensorrt_llm/_torch/disaggregation/blob/keys.py @@ -2,7 +2,7 @@ # 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 +A private encoding of the lender's opaque object keys, 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. @@ -18,16 +18,32 @@ _SEPARATOR = "/" -def encode_key(prefix: str, name: bytes) -> str: - """``"/v/"``. +def check_key_prefix(prefix: str) -> None: + """Raises ``ValueError`` unless ``prefix`` matches ``[A-Za-z0-9_-]{1,64}``.""" + if not _PREFIX.fullmatch(prefix): + raise ValueError(f"bad key prefix {prefix!r}: expected [A-Za-z0-9_-]{{1,64}}") + + +def key_stem(prefix: str) -> str: + """``"/v/"``, how every key under ``prefix`` starts. A caller encoding + many object keys checks the prefix once and appends each object key's hex to the stem. + + Args: + prefix: As for ``encode_key``. + """ + check_key_prefix(prefix) + return f"{prefix}{_SEPARATOR}v{FORMAT_VERSION}{_SEPARATOR}" + + +def encode_key(prefix: str, object_key: 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. + object_key: The lender's opaque object key; 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()}" + stem = key_stem(prefix) + if not object_key: + raise ValueError("an object key is non-empty") + return stem + object_key.hex() diff --git a/tensorrt_llm/_torch/disaggregation/blob/store.py b/tensorrt_llm/_torch/disaggregation/blob/store.py index acf2d60fede9..3f5a50091a43 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/store.py +++ b/tensorrt_llm/_torch/disaggregation/blob/store.py @@ -12,9 +12,11 @@ from __future__ import annotations import enum -from collections.abc import Sequence +from collections.abc import Callable, Iterator, Sequence from dataclasses import dataclass -from typing import Any, Protocol, runtime_checkable +from typing import Any, Protocol, overload, runtime_checkable + +import numpy as np from ..base.memory import MemoryKind @@ -22,8 +24,8 @@ """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.""" +driver it rejects the submission that carries one, before dispatch, so such a driver is never +called with a hint.""" class DriverRejected(Exception): @@ -31,8 +33,8 @@ class DriverRejected(Exception): 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. + raise from ``register``, ``put`` or ``get`` may have escaped, so the caller keeps the memory + protected; the blob backend, which cannot tell what escaped, holds it for good. """ @@ -46,7 +48,10 @@ class RegistrationError(StoreError): 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. + anyway, or ``None`` when the store certainly did not. A driver that cannot tell gives + the handle. A failure it has no handle for is not a ``RegistrationError(None)``: the + driver lets the original exception escape, and the caller keeps the region's memory + protected. """ def __init__(self, reason: str, handle: RegionHandle | None = None) -> None: @@ -78,19 +83,22 @@ class RowStatus(enum.Enum): 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.""" + nothing of its own. ``probe`` gives the same answer for such a key. Only for a hold the store + confirmed; a store that cannot tell an unfinished write from another refusal answers FAULT.""" 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.""" + """For a put: the store would not take the row, for space or policy, and said so. Nothing was + published; the other rows of the operation go on. A refusal the driver cannot place is + FAULT.""" 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.""" + """Anything else: timeout, size or integrity mismatch, transport or store error, and any + error the driver cannot confirm to be one of the states above. Says nothing about the + destination's contents.""" @dataclass(frozen=True) @@ -202,6 +210,144 @@ def nbytes(self) -> int: return sum(span.nbytes for span in self.spans) +class RowBatch(Sequence[RowIO]): + """The rows of one ``put`` or ``get``, laid out alike and held as columns. + + Row ``i`` is stored under ``keys()[i]``; its payload is its spans + ``(addresses[i, j], span_bytes[j])`` back to back, in order. ``addresses`` is + ``int64 [rows, spans_per_row]``, read only. Every row has the same span sizes, so every row + holds ``row_bytes``. Span ``j`` of every row lies in ``extents[j]``, and each extent lies in one + region its maker registered with the driver, so a driver can check a batch's memory once per + extent instead of once per row. + + A batch keeps arrays and builds nothing per row until a driver asks, so handing one over costs + the same for any number of rows. A driver whose store calls take lists builds them on the + thread that makes those calls (``keys()``, ``addresses[:, 0].tolist()``); a driver that works + row by row reads the batch as a ``Sequence[RowIO]``, one ``RowIO`` per row read. Slicing a + batch gives the batch of those rows, again without building anything per row. + + Its maker guarantees what a driver would otherwise check row by row: at least one row, no key + twice, spans of positive size inside their extent that do not overlap within a row, and each + extent inside one registered region. :meth:`of` checks all but the last, which only the driver + can; :meth:`of_slots`, for a maker that checked its rows already, checks nothing. + """ + + def __init__( + self, + keys: Sequence[str], + addresses: np.ndarray, + span_bytes: tuple[int, ...], + extents: tuple[Span, ...], + ) -> None: + self._keys = keys + self.addresses = addresses + self.span_bytes = span_bytes + self.extents = extents + + @classmethod + def of_slots( + cls, + key_stem: str, + object_keys: np.ndarray, + addresses: np.ndarray, + slot_bytes: int, + extent: Span, + ) -> RowBatch: + """Rows that are each one whole slot: row ``i`` is keyed ``key_stem`` followed by + ``object_keys[i]`` in hex, and its payload is the ``slot_bytes`` at ``addresses[i]``. + Checks nothing and copies nothing. + + Args: + key_stem: How every key starts (``keys.key_stem``). + object_keys: ``uint8 [rows, width > 0]``, no row twice. + addresses: ``int64 [rows]``, each the start of a slot inside ``extent``. + slot_bytes: Every row's size. + extent: The registered region holding every slot, or a span inside it. + """ + column = addresses[:, None] # a view of the caller's array, which stays writable + column.flags.writeable = False + return cls(_HexKeys(key_stem, object_keys), column, (slot_bytes,), (extent,)) + + @classmethod + def of(cls, rows: Sequence[RowIO]) -> RowBatch: + """The batch of ``rows``, built row by row and checked. Span ``j``'s extent is the + smallest span holding every row's span ``j``. + + Raises: + ValueError: No rows, a key twice, or rows not laid out alike. + """ + rows = list(rows) + if not rows: + raise ValueError("a batch needs at least one row") + keys = tuple(row.key for row in rows) + if len(set(keys)) != len(keys): + raise ValueError("a key appears twice in one batch") + span_bytes = tuple(span.nbytes for span in rows[0].spans) + if any(tuple(span.nbytes for span in row.spans) != span_bytes for row in rows): + raise ValueError("the rows of a batch have the same span sizes, in the same order") + addresses = np.array([[s.address for s in row.spans] for row in rows], dtype=np.int64) + addresses.flags.writeable = False + extents = tuple( + Span(int(column.min()), int(column.max()) + nbytes - int(column.min())) + for column, nbytes in zip(addresses.T, span_bytes) + ) + return cls(keys, addresses, span_bytes, extents) + + @property + def spans_per_row(self) -> int: + return len(self.span_bytes) + + @property + def row_bytes(self) -> int: + """Every row's size: its spans' bytes added up.""" + return sum(self.span_bytes) + + def keys(self) -> list[str]: + """Every row's store key, in order.""" + return list(self._keys) + + def __len__(self) -> int: + return len(self.addresses) + + @overload + def __getitem__(self, index: int) -> RowIO: ... + + @overload + def __getitem__(self, index: slice) -> RowBatch: ... + + def __getitem__(self, index): + if isinstance(index, slice): + start, stop, step = index.indices(len(self)) + if step != 1: + raise ValueError("a RowBatch slice is contiguous") + rows = slice(start, stop) + return RowBatch(self._keys[rows], self.addresses[rows], self.span_bytes, self.extents) + spans = zip(self.addresses[index].tolist(), self.span_bytes) + return RowIO(self._keys[index], tuple(Span(address, nbytes) for address, nbytes in spans)) + + +@dataclass(frozen=True) +class _HexKeys(Sequence[str]): + """Store keys as ``stem`` followed by each row of ``object_keys`` in hex, built only when + read; iterating builds every key with one hex conversion.""" + + stem: str + object_keys: np.ndarray # uint8 [rows, width] + + def __len__(self) -> int: + return len(self.object_keys) + + def __getitem__(self, index): + if isinstance(index, slice): + return _HexKeys(self.stem, self.object_keys[index]) + return self.stem + self.object_keys[index].tobytes().hex() + + def __iter__(self) -> Iterator[str]: + stem, width = self.stem, 2 * self.object_keys.shape[1] + hexed = self.object_keys.tobytes().hex() + return iter([stem + hexed[i : i + width] for i in range(0, len(hexed), width)]) + + class RegionHandle(Protocol): """Opaque handle of one registration; only the driver that issued it reads it.""" @@ -210,6 +356,18 @@ class OpHandle(Protocol): """Opaque handle of one put or get; only the driver that issued it reads it.""" +OnChange = Callable[[OpHandle], None] +"""What a caller of ``put`` or ``get`` hands the driver to hear of the operation. + +The driver calls it with the operation's handle once every row concluded, once its access ended +and once its access turned unprovable, each time after the change shows through ``poll``, +``access_ended`` and ``access_unprovable`` (one call may tell of several). The driver may also +call it at other times, such as after some rows concluded, from any thread, outside its locks, +possibly before the ``put`` or ``get`` that issued the operation returned. The callback may call +the driver's non-blocking queries and ``cancel``; it returns promptly and never raises. A caller +that passes None polls instead.""" + + class ProbeStatus(enum.Enum): """What a probe knows about one key, without moving payload.""" @@ -230,12 +388,15 @@ class DriverCapabilities: 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. + cancel: ``cancel`` does something. Without it the backend never calls ``cancel`` and + declares no cancel of its own: the calls of a stopped submission run to their end. + probe: ``probe`` answers. The blob backend never probes, so it fetches without asking + first either way. 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 + absent_is_miss: A get of a key the store does not hold concludes MISS, unless an earlier + read of that row in the same operation already wrote its memory (a reread). 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. @@ -265,10 +426,16 @@ 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. + ``put`` and ``get`` are called on the executor loop's thread, so they only record the + operation and queue its work: they never wait on the store or start a thread, take only locks + that no thread holds while waiting on the store (letting go of the GIL only while such a lock + is contended), and do no Python work per row; one list built at C speed, such as a pending + result per row, is the most. A driver whose store calls block runs them on threads of its own, + and it alone decides how they are batched and in what order gets and puts run. + ``put``, ``get``, ``cancel`` and ``probe`` may run concurrently with each other and with the + calls above. ``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 @@ -321,22 +488,26 @@ def unregister(self, handle: RegionHandle) -> None: """ ... - def put(self, rows: Sequence[RowIO]) -> OpHandle: - """Publish ``rows``; may return before any byte moved. + def put(self, rows: RowBatch, *, on_change: OnChange | None = None) -> OpHandle: + """Publish ``rows``; returns without waiting on the store, usually before any byte moved. + ``on_change`` is called as ``OnChange`` says. 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. + ``access_ended`` is true, and hands rows of more than one span only to a driver that + declares ``multi_span``; a driver may treat unexpected ones 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. + DriverRejected: ``put`` is not declared, the driver is draining or closed, ``rows`` + exceed ``max_rows_per_op``, or an extent lies outside every 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. + def get( + self, rows: RowBatch, hint: Hint = None, *, on_change: OnChange | None = None + ) -> OpHandle: + """Fetch ``rows``, from where ``hint`` says when given; returns without waiting on the + store, usually before any byte moved. ``on_change`` is called as ``OnChange`` says. 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 @@ -395,7 +566,10 @@ def drain(self, budget_s: float) -> bool: 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. + the result. Before it returns, every row the driver has not started working on (no + transfer or store call taken up) concludes CANCELLED, and an operation left with nothing + running ends its access. 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. """ @@ -408,3 +582,59 @@ def close(self) -> None: otherwise raises :class:`DriverRejected` and changes nothing. Idempotent. """ ... + + +@dataclass(frozen=True) +class RegionSpec: + """One region of a batch registration: the arguments of ``StoreDriver.register`` as a value.""" + + address: int + nbytes: int + layout: RegionLayout + + def __post_init__(self) -> None: + Span(self.address, self.nbytes) # the same checks, the same ValueError + if self.layout.nbytes > self.nbytes: + raise ValueError(f"layout needs {self.layout.nbytes} bytes, region has {self.nbytes}") + + @property + def span(self) -> Span: + return Span(self.address, self.nbytes) + + +@runtime_checkable +class RegistersRegionsTogether(Protocol): + """Optional: a driver whose store takes a set of regions in one step and lets them go in one + step, as one that opens its client with the memory it will use. A driver declares it by + implementing it; the backend checks with ``isinstance`` and then registers a lender's regions + through it instead of one ``register`` per region. + + No atomicity is promised. ``register_regions`` fails as ``register`` does: ``DriverRejected`` + or ``RegistrationError(handle=None)`` say that nothing of the batch is registered; + ``RegistrationError(handle=h)`` says that some of it may be and ``h`` revokes whatever is; any + other raise says that some of it may be registered with nothing to revoke it by, and the + caller keeps the memory protected. + + Whether such a driver also takes ``register`` and ``unregister`` for single regions is its own + business; one that does not rejects them with ``DriverRejected``. Threads: as ``register`` and + ``unregister``, one call at a time among the four. + """ + + def register_regions(self, specs: Sequence[RegionSpec]) -> RegionHandle: + """Register every region of ``specs`` and return one handle for the batch. + + Raises: + DriverRejected: Bad geometry, a region overlapping another of the batch or a live + registration, or a DEVICE region for a driver without ``device_memory``; nothing + escaped. + RegistrationError: The store refused; ``handle`` is set when some of the batch may be + registered anyway. + """ + ... + + def unregister_regions(self, handle: RegionHandle) -> None: + """Revoke the batch, under the rules of ``unregister``: only once every operation touching + its regions has ``access_ended``, otherwise :class:`DriverRejected` and the batch stays + live; idempotent once it succeeded; any other raise leaves it live and the call + retryable.""" + ... diff --git a/tensorrt_llm/_torch/disaggregation/testing/__init__.py b/tensorrt_llm/_torch/disaggregation/testing/__init__.py index 6ad7306bcf28..39cedb419577 100644 --- a/tensorrt_llm/_torch/disaggregation/testing/__init__.py +++ b/tensorrt_llm/_torch/disaggregation/testing/__init__.py @@ -3,9 +3,10 @@ """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. +them; ``memory_store`` holds the in-process store and the drivers over it that stand in for a +real store, one registering regions singly and one in batches; ``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 diff --git a/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py b/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py index 69235c6e2b1d..78a66fcc4b19 100644 --- a/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py +++ b/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py @@ -10,6 +10,7 @@ from __future__ import annotations import ctypes +import threading import time from collections.abc import Sequence from typing import Any @@ -18,10 +19,14 @@ DriverCapabilities, DriverRejected, Hint, + OnChange, OpHandle, ProbeStatus, RegionHandle, RegionLayout, + RegionSpec, + RegistrationError, + RowBatch, RowIO, RowResult, RowStatus, @@ -30,24 +35,38 @@ faulted, ) -from .memory_store import DEFAULT_CAPABILITIES, Kind, MemoryDriver, MemoryOp, MemoryStore, _scatter +from .memory_store import ( + DEFAULT_CAPABILITIES, + Kind, + MemoryDriver, + MemoryOp, + MemoryStore, + ObjectNotReady, + TogetherMemoryDriver, + _Batch, + _cancellable, + _Registration, + _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: + def get( + self, rows: RowBatch, hint: Hint = None, *, on_change: OnChange | None = None + ) -> MemoryOp: for io in rows: for span in io.spans: ctypes.memset(span.address, 0, span.nbytes) - return super().get(rows, hint) + return super().get(rows, hint, on_change=on_change) + + +class ChecksRowsNotExtents(MemoryDriver): + def _check_extents(self, rows: Sequence[RowIO]) -> None: + pass class ChecksOnlyTheStartOfASpan(MemoryDriver): @@ -67,7 +86,10 @@ def _get_row(self, io: RowIO) -> RowResult: class TruncatesLargerObjects(MemoryDriver): def _get_row(self, io: RowIO) -> RowResult: - stored = self.store.get(io.key) + try: + stored = self.store.get(io.key) + except ObjectNotReady: # another writer's: a sound get answers that + return super()._get_row(io) if stored is not None and len(stored.payload) > io.nbytes: _scatter(io.spans, stored.payload[: io.nbytes]) return _READY @@ -76,7 +98,10 @@ def _get_row(self, io: RowIO) -> RowResult: class TruncatesIntoMultiSpanRows(MemoryDriver): def _get_row(self, io: RowIO) -> RowResult: - stored = self.store.get(io.key) + try: + stored = self.store.get(io.key) + except ObjectNotReady: # another writer's: a sound get answers that + return super()._get_row(io) 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 @@ -106,18 +131,20 @@ def _put_row(self, io: RowIO) -> RowResult: 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)) + if self.store.presence(io.key) is ProbeStatus.HIT: # only an object that exists + 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): + """Clears the whole region a get row lies in, then delivers the row itself as it should.""" + 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 + return super()._get_row(io) class OrdersSpansByAddress(MemoryDriver): @@ -128,10 +155,161 @@ def _get_row(self, io: RowIO) -> RowResult: return super()._get_row(RowIO(io.key, tuple(sorted(io.spans, key=lambda s: s.address)))) +class ChecksOnlyTheFirstSpan(MemoryDriver): + def _check_registered(self, rows: Sequence[RowIO]) -> None: + super()._check_registered([RowIO(io.key, io.spans[:1]) for io in rows]) + + def _check_extents(self, rows: Sequence[RowIO]) -> None: + pass # nor the extents, which would catch the lapse above + + +class AcceptsOverlappingRegistrations(MemoryDriver): + def register(self, address: int, nbytes: int, layout: RegionLayout) -> RegionHandle: + with self._lock: + registration = _Registration(Span(address, nbytes), layout) + self._registrations.append(registration) + return registration + + +class IgnoresUnregistration(MemoryDriver): + def _registered(self, span: Span) -> bool: + return any(span.within(r.span) for r in self._registrations) + + +class SecondUnregisterRaises(MemoryDriver): + def unregister(self, handle: RegionHandle) -> None: + with self._lock: + if not self._registration(handle).live: + raise DriverRejected("already unregistered") + super().unregister(handle) + + +class UnregistersTheRegionOfAnUnprovableOperation(MemoryDriver): + def _check_region_idle(self, registration: _Registration) -> None: + if any( + not op.access_ended + and not op.unprovable + and any(span.within(registration.span) for row in op.rows for span in row.io.spans) + for op in self._ops + ): + raise DriverRejected("an operation may still access the region") + + +class DuplicatePutIsAFault(MemoryDriver): + def _put_row(self, io: RowIO) -> RowResult: + if self.store.presence(io.key) is ProbeStatus.HIT: + return faulted("already exists") + return super()._put_row(io) + + +class CallsAHeldPutReady(MemoryDriver): + def _put_row(self, io: RowIO) -> RowResult: + result = super()._put_row(io) + return _READY if result.status is RowStatus.NOT_READY else result + + +class ProbeCallsAHeldKeyAMiss(MemoryDriver): + def probe(self, keys: Sequence[str], budget_s: float, hint: Hint = None): + answers = super().probe(keys, budget_s, hint) + return tuple(ProbeStatus.MISS if a is ProbeStatus.NOT_READY else a for a in answers) + + +class ProbeCallsAMissAHit(MemoryDriver): + def probe(self, keys: Sequence[str], budget_s: float, hint: Hint = None): + answers = super().probe(keys, budget_s, hint) + return tuple(ProbeStatus.HIT if a is ProbeStatus.MISS else a for a in answers) + + +class SplitsAMultiSpanRowEvenly(MemoryDriver): + """Scatters a multi-span get as if every span had the same size.""" + + def _get_row(self, io: RowIO) -> RowResult: + if len(io.spans) == 1: + return super()._get_row(io) + stored = self.store.get(io.key) + if stored is not None and len(stored.payload) == io.nbytes: + share = io.nbytes // len(io.spans) + for k, span in enumerate(io.spans): + ctypes.memmove( + span.address, stored.payload[k * share : k * share + span.nbytes], span.nbytes + ) + return _READY + return super()._get_row(io) + + +class ResultsChangeAfterTheyConcluded(MemoryDriver): + def poll(self, op: OpHandle) -> tuple[RowResult, ...]: + results = super().poll(op) + polls = self.__dict__.setdefault("_polls", {}) + polls[id(op)] = polls.get(id(op), 0) + 1 + if polls[id(op)] >= 3 and all(r.concluded for r in results): + return tuple(faulted("changed") if r.status is RowStatus.READY else r for r in results) + return results + + +class CancelRewritesConcludedRows(MemoryDriver): + def _cancel_unstarted(self, op: MemoryOp) -> None: + for row in op.rows: + row.published = _CANCELLED + self._settle_untouched(op) + + +class CancelSkipsOnChange(MemoryDriver): + def cancel(self, op: OpHandle) -> None: + with self._lock: + if self._cancel_effective: + self._cancel_unstarted(self._op(op)) + + +class DrainSkipsOnChange(MemoryDriver): + def drain(self, budget_s: float) -> bool: + with self._lock: + self._require_not_closed() + self._start_draining() + for op in self._ops: + if any(_cancellable(row) for row in op.rows): + self._cancel_unstarted(op) + self._lock.wait_for(self._drain_can_return, timeout=budget_s) + return self._quiet() + + +class DrainForcesAnEnd(MemoryDriver): + """Once its budget ran out, declares the access of every operation still running ended.""" + + def drain(self, budget_s: float) -> bool: + drained = super().drain(budget_s) + if not drained: + with self._lock: + for op in self._ops: + op.access_ended = op.access_ended or not op.unprovable + return drained + + +class GetLetsACrashEscape(MemoryDriver): + def _move(self, kind: Kind, io: RowIO) -> RowResult: + if kind != "get": + return super()._move(kind, io) + try: + return self._get_row(io) + except StoreError as error: # only what is not a store error escapes + return faulted(str(error)) + + +class GetWaitsForTheStore(MemoryDriver): + """Runs a get on a thread of its own, but returns its handle only once the store answered.""" + + def get( + self, rows: RowBatch, hint: Hint = None, *, on_change: OnChange | None = None + ) -> MemoryOp: + op = super().get(rows, hint, on_change=on_change) + while not all(row.published.concluded for row in op.rows): + time.sleep(0.001) + return op + + class IgnoresTheRowLimit(MemoryDriver): def _check_row_limit(self, rows: Sequence[RowIO]) -> None: - if not rows: - raise DriverRejected("an operation needs at least one row") + pass class IgnoresTheDirection(MemoryDriver): @@ -151,6 +329,33 @@ def _execute(self, op: MemoryOp) -> None: op.access_ended = False # the evidence is lost again +class CallsOnChangeBeforeTheChangeShows(MemoryDriver): + def _run(self, op: MemoryOp, finished: threading.Event | None) -> None: + self._call_on_change(op) # nothing has changed yet, and on_change is not called again + try: + self._before_execute(op) + self._execute(op) + finally: + if finished is not None: + finished.set() + with self._lock: + self._lock.notify_all() + + +class SkipsOnChangeWhenUnprovable(MemoryDriver): + def mark_unprovable(self, op: MemoryOp) -> None: + with self._lock: + self._require_open(op) + op.unprovable = True + self._lock.notify_all() + + +class SkipsOnChangeWhileAccessIsOpen(MemoryDriver): + def _call_on_change(self, op: MemoryOp) -> None: + if op.access_ended or op.unprovable: + super()._call_on_change(op) + + class ChangesItsCapabilities(MemoryDriver): _flip = False @@ -245,6 +450,16 @@ def unregister(self, handle: RegionHandle) -> None: pass +class PutWaitsForTheStore(MemoryDriver): + """Runs a put on a thread of its own, but returns its handle only once the store answered.""" + + def put(self, rows: RowBatch, *, on_change: OnChange | None = None) -> MemoryOp: + op = super().put(rows, on_change=on_change) + while not all(row.published.concluded for row in op.rows): + time.sleep(0.001) + return op + + class CancelWaitsForTheTransfer(MemoryDriver): def cancel(self, op: OpHandle) -> None: with self._lock: @@ -260,6 +475,13 @@ def _cancel_unstarted(self, op: MemoryOp) -> None: self._settle_untouched(op) +class EndsAccessOnceARowStarted(MemoryDriver): + def access_ended(self, op: OpHandle) -> bool: + with self._lock: + recorded = self._op(op) + return recorded.access_ended or any(row.started for row in recorded.rows) + + 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.""" @@ -271,12 +493,15 @@ def __init__(self, store: MemoryStore, **options: Any) -> None: def cancel(self, op: OpHandle) -> None: with self._lock: recorded = self._op(op) + if self._cancel_effective: + self._cancel_unstarted(recorded) # as a sound cancel does 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() + self._call_on_change(recorded) def _move(self, kind: Kind, io: RowIO) -> RowResult: result = super()._move(kind, io) @@ -289,6 +514,9 @@ def _move(self, kind: Kind, io: RowIO) -> RowResult: GET_ONLY = DriverCapabilities( get=True, put=False, cancel=True, probe=True, max_rows_per_op=None, absent_is_miss=True ) +PUT_ONLY = DriverCapabilities( + get=False, put=True, cancel=True, probe=True, max_rows_per_op=None, absent_is_miss=True +) HOST_ONLY = DriverCapabilities( get=True, put=True, @@ -301,16 +529,16 @@ def _move(self, kind: Kind, io: RowIO) -> RowResult: # 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 batch's extents are checked, not only its rows": ( + ChecksRowsNotExtents, + DEFAULT_CAPABILITIES, + "test_a_batch_whose_extent_leaves_its_registration_is_rejected_before_any_effect", + ), "a span is checked to its end": ( ChecksOnlyTheStartOfASpan, DEFAULT_CAPABILITIES, @@ -381,6 +609,21 @@ def _move(self, kind: Kind, io: RowIO) -> RowResult: DEFAULT_CAPABILITIES, "test_results_are_stable_and_access_ends_with_evidence", ), + "on_change follows the change it tells of": ( + CallsOnChangeBeforeTheChangeShows, + DEFAULT_CAPABILITIES, + "test_on_change_is_called_after_each_operation_settles", + ), + "on_change tells of rows concluded before access ends": ( + SkipsOnChangeWhileAccessIsOpen, + DEFAULT_CAPABILITIES, + "test_on_change_is_called_once_every_row_concluded_while_access_stays_open", + ), + "on_change tells of an access turned unprovable": ( + SkipsOnChangeWhenUnprovable, + DEFAULT_CAPABILITIES, + "test_an_unprovable_get_calls_on_change_keeps_its_region_and_drain_returns_at_once", + ), "capabilities are fixed": ( ChangesItsCapabilities, DEFAULT_CAPABILITIES, @@ -404,12 +647,12 @@ def _move(self, kind: Kind, io: RowIO) -> RowResult: "a blocked get delays nothing else": ( HoldsTheLockAcrossTheStore, DEFAULT_CAPABILITIES, - "test_a_blocked_get_delays_neither_another_get_nor_the_non_blocking_calls", + "test_a_get_blocked_in_the_store_delays_nothing_and_keeps_its_memory", ), "a busy region is not unregistered": ( UnregistersBusyRegions, DEFAULT_CAPABILITIES, - "test_a_blocked_get_delays_neither_another_get_nor_the_non_blocking_calls", + "test_a_get_blocked_in_the_store_delays_nothing_and_keeps_its_memory", ), "an unfinished write is not ready": ( CallsNotReadyAFault, @@ -423,13 +666,13 @@ def _move(self, kind: Kind, io: RowIO) -> RowResult: ), "a recorded operation never raises": ( LetsRaisesEscape, - DEFAULT_CAPABILITIES, + PUT_ONLY, # a one-direction driver: the check must run, not skip "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", + "test_an_unprovable_get_calls_on_change_keeps_its_region_and_drain_returns_at_once", ), "drain is one-way": ( KeepsAcceptingWhileDraining, @@ -444,13 +687,83 @@ def _move(self, kind: Kind, io: RowIO) -> RowResult: "unregister frees the range": ( UnregisterKeepsTheRange, DEFAULT_CAPABILITIES, - "test_unregister_frees_the_range_for_a_new_registration", + "test_a_registered_range_cannot_be_registered_again_until_unregistered", + ), + "registrations never overlap": ( + AcceptsOverlappingRegistrations, + DEFAULT_CAPABILITIES, + "test_a_registered_range_cannot_be_registered_again_until_unregistered", + ), + "every span is checked": ( + ChecksOnlyTheFirstSpan, + DEFAULT_CAPABILITIES, + "test_every_span_of_a_row_is_checked_against_the_registrations", + ), + "an unregistered span is rejected": ( + IgnoresUnregistration, + DEFAULT_CAPABILITIES, + "test_a_span_leaving_or_outliving_its_registration_is_rejected", + ), + "unregister is idempotent": ( + SecondUnregisterRaises, + DEFAULT_CAPABILITIES, + "test_unregister_is_idempotent", + ), + "an unprovable operation keeps its region": ( + UnregistersTheRegionOfAnUnprovableOperation, + DEFAULT_CAPABILITIES, + "test_an_unprovable_get_calls_on_change_keeps_its_region_and_drain_returns_at_once", + ), + "a duplicate put is ready": ( + DuplicatePutIsAFault, + DEFAULT_CAPABILITIES, + "test_a_duplicate_put_of_a_committed_object_is_ready", + ), + "a put of a held key is not ready": ( + CallsAHeldPutReady, + DEFAULT_CAPABILITIES, + "test_a_key_held_by_an_unfinished_write_is_not_ready_for_get_put_and_probe", + ), + "a probe of a held key is not ready": ( + ProbeCallsAHeldKeyAMiss, + DEFAULT_CAPABILITIES, + "test_a_key_held_by_an_unfinished_write_is_not_ready_for_get_put_and_probe", + ), + "probe tells hit from miss": ( + ProbeCallsAMissAHit, + DEFAULT_CAPABILITIES, + "test_probe_reports_hit_and_miss", + ), + "spans of a row keep their sizes": ( + SplitsAMultiSpanRowEvenly, + DEFAULT_CAPABILITIES, + "test_a_row_may_cover_parts_of_slots_across_regions", + ), + "results never change once concluded": ( + ResultsChangeAfterTheyConcluded, + DEFAULT_CAPABILITIES, + "test_results_are_stable_and_access_ends_with_evidence", + ), + "cancel leaves concluded rows alone": ( + CancelRewritesConcludedRows, + DEFAULT_CAPABILITIES, + "test_cancel_after_conclusion_changes_nothing", + ), + "a recorded get never raises": ( + GetLetsACrashEscape, + GET_ONLY, # a one-direction driver: the check must run, not skip + "test_a_raise_after_the_operation_is_recorded_becomes_a_row_fault", ), } # 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]] = { + "put never waits for the store": ( + PutWaitsForTheStore, + DEFAULT_CAPABILITIES, + "test_a_put_returns_its_handle_before_the_store_answers", + ), "cancel never waits for the store": ( CancelWaitsForTheTransfer, DEFAULT_CAPABILITIES, @@ -466,4 +779,160 @@ def _move(self, kind: Kind, io: RowIO) -> RowResult: DEFAULT_CAPABILITIES, "test_cancel_of_a_get_inside_the_store_neither_waits_nor_lies", ), + "access stays open while a row is in the store": ( + EndsAccessOnceARowStarted, + DEFAULT_CAPABILITIES, + "test_a_get_blocked_in_the_store_delays_nothing_and_keeps_its_memory", + ), + "get never waits for the store": ( + GetWaitsForTheStore, + DEFAULT_CAPABILITIES, + "test_a_get_blocked_in_the_store_delays_nothing_and_keeps_its_memory", + ), + "drain never forces an end": ( + DrainForcesAnEnd, + DEFAULT_CAPABILITIES, + "test_a_get_blocked_in_the_store_delays_nothing_and_keeps_its_memory", + ), + "a cancel tells of the operation it ended": ( + CancelSkipsOnChange, + DEFAULT_CAPABILITIES, + "test_a_cancel_that_ends_an_operation_calls_its_on_change", + ), + "a drain tells of the operations it ended": ( + DrainSkipsOnChange, + DEFAULT_CAPABILITIES, + "test_a_drain_that_ends_an_operation_calls_its_on_change", + ), +} + + +# Drivers breaking a rule of the optional RegistersRegionsTogether protocol. + + +class KeepsTheBatchAfterAFailure(TogetherMemoryDriver): + """A failed ``register_regions`` that certifies nothing is registered while the store kept + the batch: the memory is now registered with nothing to revoke it by.""" + + def register_regions(self, specs: Sequence[RegionSpec]) -> _Batch: + try: + return super().register_regions(specs) + except RegistrationError as error: + if error.handle is None: + with self._lock: + self._keep(_Batch([_Registration(spec.span, spec.layout) for spec in specs])) + raise + + +class RevokesABatchOnlyOnce(TogetherMemoryDriver): + def unregister_regions(self, handle: RegionHandle) -> None: + with self._lock: + if not self._batch(handle).live: + raise StoreError("the batch was already revoked") + super().unregister_regions(handle) + + +class RevocationKeepsTheRange(TogetherMemoryDriver): + def unregister_regions(self, handle: RegionHandle) -> None: + with self._lock: + batch = self._batch(handle) + for registration in batch.registrations: + self._check_region_idle(registration) + batch.live = False # its registrations stay live: the range is never freed + + +class RegistersOnlyTheFirstRegion(TogetherMemoryDriver): + def register_regions(self, specs: Sequence[RegionSpec]) -> _Batch: + return super().register_regions(specs[:1]) + + +class RevokesBusyBatches(TogetherMemoryDriver): + def _check_region_idle(self, registration: _Registration) -> None: + pass + + +class RevokesTheIdleRegionsOfABusyBatch(TogetherMemoryDriver): + """Refuses a busy batch, but has already let its idle regions go.""" + + def unregister_regions(self, handle: RegionHandle) -> None: + with self._lock: + busy = 0 + for registration in self._batch(handle).registrations: + try: + self._check_region_idle(registration) + except DriverRejected: + busy += 1 + else: + registration.live = False + if busy: + raise DriverRejected(f"{busy} region(s) of the batch are busy") + super().unregister_regions(handle) + + +class RevokesTheIdleRegionsOfABusyBatchOneAtATime(RevokesTheIdleRegionsOfABusyBatch): + """The same, in a driver that takes one batch at a time: it refuses a new registration + whatever the idle regions' state, so only a row moved through them tells they are gone.""" + + def register_regions(self, specs: Sequence[RegionSpec]) -> _Batch: + with self._lock: + if self.live_batches(): + raise DriverRejected("one batch at a time") + return super().register_regions(specs) + + +class UnregisterRegionsWaitsForTheTransfer(TogetherMemoryDriver): + """Revokes a busy batch once it is idle, instead of refusing: a caller holding something the + transfer needs is stuck.""" + + def unregister_regions(self, handle: RegionHandle) -> None: + with self._lock: + self._lock.wait_for(lambda: all(op.access_ended for op in self._ops)) + super().unregister_regions(handle) + + +# The same, for the checks of RegistersRegionsTogether; the pairing test runs these against the +# in-memory driver that registers in batches. +BROKEN_TOGETHER_CONFORMANCE: dict[ + str, tuple[type[TogetherMemoryDriver], DriverCapabilities, str] +] = { + "a failed batch certifies only what is true": ( + KeepsTheBatchAfterAFailure, + DEFAULT_CAPABILITIES, + "test_a_failed_batch_registration_is_revocable_or_left_nothing_behind", + ), + "unregister_regions is idempotent": ( + RevokesABatchOnlyOnce, + DEFAULT_CAPABILITIES, + "test_regions_registered_together_are_revoked_together_and_the_range_freed", + ), + "unregister_regions frees the range": ( + RevocationKeepsTheRange, + DEFAULT_CAPABILITIES, + "test_regions_registered_together_are_revoked_together_and_the_range_freed", + ), + "every region of a batch is registered": ( + RegistersOnlyTheFirstRegion, + DEFAULT_CAPABILITIES, + "test_regions_registered_together_are_revoked_together_and_the_range_freed", + ), + "a busy batch is not revoked": ( + RevokesBusyBatches, + DEFAULT_CAPABILITIES, + "test_a_batch_is_not_revoked_while_any_of_its_regions_is_busy", + ), + "a busy batch keeps its idle regions": ( + RevokesTheIdleRegionsOfABusyBatch, + DEFAULT_CAPABILITIES, + "test_a_batch_is_not_revoked_while_any_of_its_regions_is_busy", + ), + "a busy batch keeps its idle regions, one batch at a time": ( + RevokesTheIdleRegionsOfABusyBatchOneAtATime, + DEFAULT_CAPABILITIES, + "test_a_batch_is_not_revoked_while_any_of_its_regions_is_busy", + ), + "unregister_regions never waits for the store": ( + UnregisterRegionsWaitsForTheTransfer, + DEFAULT_CAPABILITIES, + "test_a_batch_is_not_revoked_while_any_of_its_regions_is_busy", + ), } diff --git a/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py b/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py index a5d672296024..64beb2d1aaae 100644 --- a/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py +++ b/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py @@ -24,19 +24,25 @@ import threading import time from collections.abc import Callable, Generator, Sequence +from typing import NoReturn +import numpy as np 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.keys import encode_key, key_stem from tensorrt_llm._torch.disaggregation.blob.store import ( DriverRejected, Hint, + OnChange, OpHandle, ProbeStatus, RegionHandle, RegionLayout, + RegionSpec, + RegistersRegionsTogether, RegistrationError, + RowBatch, RowIO, RowResult, RowStatus, @@ -122,8 +128,14 @@ def register(self) -> None: 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 rows(self) -> RowBatch: + """Every label's row, one whole slot each, laid out as the blob backend hands rows over: + keys built only when read, and addresses as a column of the caller's array.""" + names = np.frombuffer(b"".join(name_of(label) for label in self.labels), dtype=np.uint8) + slots = self.memory.address + self.slot * np.arange(len(self.labels), dtype=np.int64) + region = Span(self.memory.address, self.memory.nbytes) + object_keys = names.reshape(len(self.labels), NAME_BYTES) + return RowBatch.of_slots(key_stem(self.prefix), object_keys, slots, self.slot, region) def span_at(self, offset: int, nbytes: int) -> Span: """Any range of the region, slot boundaries or not.""" @@ -152,6 +164,65 @@ def unregister(self) -> None: # Never freed. UNRELEASED_ENDPOINTS: list[Endpoint] = [] +# Memory a batch registration may still hold when a check ends. Never freed. +UNRELEASED_MEMORY: list[Memory] = [] + + +def _revoke_again(driver: RegistersRegionsTogether, handle: RegionHandle) -> None: + """A second ``unregister_regions`` of a revoked batch; any raise breaks idempotency.""" + try: + driver.unregister_regions(handle) + except Exception as error: # noqa: BLE001 - the rule allows no raise at all + pytest.fail(f"a second revocation of the batch raised {error!r}") + + +def _prove_registered(driver: StoreDriver, memory: Memory, index: int, timeout: float) -> bool: + """A row inside memory moves: the region is registered, as the driver refuses a row outside + one. A probe by registration could not tell, since a driver that takes one batch at a time + refuses a new one whatever the region's state. Returns whether the row's access ended: a + get-only driver without ``absent_is_miss`` may conclude the get of the fresh key as FAULT and + leave its access unprovable, after which the batch cannot be revoked.""" + row = RowIO(key_of(f"R{index}", "t" + secrets.token_hex(4)), (Span(memory.address, SLOT),)) + try: + if driver.capabilities.put: + memory.write(b"a" * SLOT, 0) + op = driver.put(RowBatch.of([row])) + else: + op = driver.get(RowBatch.of([row])) + except DriverRejected as error: + pytest.fail(f"region {index} of the batch takes no row: {error}") + wait_until(lambda: concluded(driver.poll(op)), "the operation to conclude", timeout) + if driver.capabilities.put: + assert statuses(driver, op) == [RowStatus.READY] + return _access_settled(driver, op, timeout) + + +def _access_settled(driver: StoreDriver, op: OpHandle, timeout: float) -> bool: + """Waits until the operation's access ended or the driver declared that it can never prove + it did; returns whether it ended.""" + wait_until( + lambda: driver.access_ended(op) or driver.access_unprovable(op), + "access to end or to be declared unprovable", + timeout, + ) + return driver.access_ended(op) + + +def _skip_unprovable(index: int) -> NoReturn: + """The contract allows a store to leave a row's access unprovable, and a batch with such a row + cannot be revoked: the check ends here, and its ``finally`` retains the memory.""" + pytest.skip( + f"region {index}: the store left a row's access unprovable, so the batch stays registered " + "and its memory is retained" + ) + + +def specs_of(memory: Sequence[Memory], slot: int = SLOT) -> list[RegionSpec]: + """One region per memory, laid out in slots of ``slot`` bytes, for ``register_regions``.""" + return [ + RegionSpec(m.address, m.nbytes, RegionLayout(slot, m.nbytes // slot, "kv")) for m in memory + ] + class Harness: """One driver under test plus everything a check created with it.""" @@ -193,14 +264,20 @@ def endpoint( endpoint.register() return endpoint - def put(self, rows: Sequence[RowIO]) -> OpHandle: - op = self.driver.put(rows) + def put(self, rows: Sequence[RowIO], on_change: OnChange | None = None) -> OpHandle: + op = self.driver.put(_batch(rows), on_change=on_change) 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) + def get( + self, rows: Sequence[RowIO], hint: Hint = None, on_change: OnChange | None = None + ) -> OpHandle: + # Without a hint the driver is called without one, as the backend calls it. + batch = _batch(rows) + if hint is None: + op = self.driver.get(batch, on_change=on_change) + else: + op = self.driver.get(batch, hint, on_change=on_change) self.ops.append(op) return op @@ -211,8 +288,8 @@ def publish(self, rows: dict[str, bytes], slot: int = SLOT) -> Endpoint: 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) + assert [r.status for r in self.wait_for_results(op)] == [RowStatus.READY] * len(rows) + self.wait_for_access_end(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() @@ -220,14 +297,14 @@ def publish(self, rows: dict[str, bytes], slot: int = SLOT) -> Endpoint: # Waiting and asserting. - def settle(self, op: OpHandle) -> tuple[RowResult, ...]: + def wait_for_results(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: + def wait_for_access_end(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: @@ -350,7 +427,14 @@ class DriverConformance: 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 + # The driver implements RegistersRegionsTogether and, with registers_on_demand False, the + # handle allocate hands out is the batch its memory was registered as; the suite then revokes + # through unregister_regions, never through unregister. + registers_together = False + # True when put/get return their handle before the store answers, as the driver protocol + # asks. A driver that answers inside put/get, as the in-memory driver's automatic mode does, + # sets False to skip the checks that hold a put or a get inside the store. + returns_before_the_store_answers = True def make_driver(self) -> StoreDriver: raise NotImplementedError @@ -389,6 +473,11 @@ def make_blocked_get(self) -> tuple[StoreDriver, str, threading.Event | None, Un runs.""" return None + def make_blocked_put(self) -> tuple[StoreDriver, str, threading.Event | None, Unblock] | None: + """A driver whose store stalls every put of the returned key until the callable runs, + with an event set once a put is stalled, or ``None`` for no such signal.""" + return None + def make_held_key(self) -> tuple[StoreDriver, str, Unblock] | None: """A driver whose store holds the returned key as another writer's unfinished object until the callable runs.""" @@ -399,7 +488,20 @@ def make_full_store(self) -> StoreDriver | None: 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.""" + """A driver whose store raises something that is not a store error from inside every put + and every get.""" + return None + + def make_driver_leaving_operations_unstarted(self) -> StoreDriver | None: + """A driver that starts no operation it takes: every row stays unstarted until a cancel or + a drain ends it.""" + return None + + def make_results_before_access_end( + self, + ) -> tuple[StoreDriver, Callable[[OpHandle], None], Callable[[OpHandle], None]] | None: + """A driver whose operations wait for the check: the first callable concludes every row of + an operation and leaves its access open, the second ends its access.""" return None def make_failing_presence(self) -> tuple[StoreDriver, str] | None: @@ -415,6 +517,25 @@ def make_unprovable_get(self) -> tuple[StoreDriver, Callable[[OpHandle], None]] it while the driver can never prove that its access ended.""" return None + def make_together_registration(self) -> tuple[StoreDriver, list[Memory]] | None: + """A fresh driver implementing ``RegistersRegionsTogether`` with no batch live, and + unregistered memory for one batch of at least two regions.""" + return None + + def make_failing_together_registrations(self) -> list[tuple[StoreDriver, list[Memory]]]: + """Fresh drivers whose next ``register_regions`` of the returned memory fails after the + store was asked, one per failure shape the store can stage; empty when it can stage + none.""" + return [] + + def make_busy_together_registration( + self, + ) -> tuple[StoreDriver, list[Memory], str, threading.Event, Unblock] | None: + """As ``make_together_registration``, plus a key whose get stalls in the store until the + callable runs, and an event set once a get of it is stalled. The event is required: the + check decides nothing by waiting; a suite that cannot give the signal returns ``None``.""" + 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 @@ -437,7 +558,7 @@ def test_put_then_get_moves_whole_rows_and_never_aliases_the_source(self) -> Non 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())) + results = h.wait_for_results(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" @@ -448,7 +569,7 @@ def test_get_reports_hit_hit_miss_and_leaves_the_miss_untouched(self) -> None: 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())) + results = h.wait_for_results(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() @@ -458,7 +579,7 @@ def test_an_object_larger_than_its_row_is_a_fault_not_a_truncated_ready(self) -> 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())) + (result,) = h.wait_for_results(h.get(reader.rows())) assert result.status is RowStatus.FAULT, result assert reader.intact() # nothing lands outside the row @@ -469,7 +590,7 @@ def test_an_object_smaller_than_its_row_is_a_fault(self) -> None: 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())) + (result,) = h.wait_for_results(h.get(reader.rows())) assert result.status is RowStatus.FAULT, result assert reader.intact() @@ -479,7 +600,7 @@ def test_a_duplicate_put_of_a_committed_object_is_ready(self) -> None: 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())) + (result,) = h.wait_for_results(h.get(reader.rows())) assert result.status is RowStatus.READY assert reader.read_row(0) == b"aaaaaaaa" @@ -491,10 +612,10 @@ def test_a_put_under_an_existing_key_leaves_the_object_whole(self) -> None: h.publish({"A": b"aaaaaaaa"}) other = h.endpoint(["A"]) other.fill_row(0, b"zzzzzzzz") - (result,) = h.settle(h.put(other.rows())) + (result,) = h.wait_for_results(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())) + (got,) = h.wait_for_results(h.get(reader.rows())) assert got.status is RowStatus.READY assert reader.read_row(0) == b"aaaaaaaa" @@ -506,9 +627,9 @@ def test_an_operation_leaves_the_rest_of_its_region_alone(self) -> None: h.publish({"B": b"bbbbbbbb"}) reader = h.endpoint(["A", "B", "C"], fill=0x2D) op = h.get([reader.rows()[1]]) - (result,) = h.settle(op) + (result,) = h.wait_for_results(op) assert result.status is RowStatus.READY - h.settle_access(op) + h.wait_for_access_end(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() @@ -521,11 +642,11 @@ def test_a_multi_span_row_is_stored_and_fetched_back_to_back(self) -> None: 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])) + (put,) = h.wait_for_results(h.put([row])) assert put.status is RowStatus.READY whole = h.endpoint(["M"], slot=8, fill=0x2D) - (got,) = h.settle(h.get(whole.rows())) + (got,) = h.wait_for_results(h.get(whole.rows())) assert got.status is RowStatus.READY assert whole.read_row(0) == b"aaaacccc" @@ -536,7 +657,7 @@ def test_a_multi_span_row_is_stored_and_fetched_back_to_back(self) -> None: ) 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])) + (got,) = h.wait_for_results(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"----" @@ -552,17 +673,17 @@ def test_a_row_may_cover_parts_of_slots_across_regions(self) -> None: 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])) + (put,) = h.wait_for_results(h.put([row])) assert put.status is RowStatus.READY whole = h.endpoint(["P"], slot=8, fill=0x2D) - (got,) = h.settle(h.get(whole.rows())) + (got,) = h.wait_for_results(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])) + (got,) = h.wait_for_results(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() @@ -573,7 +694,7 @@ def test_a_larger_object_is_not_truncated_into_a_multi_span_row(self) -> None: 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])) + (got,) = h.wait_for_results(h.get([row])) assert got.status is RowStatus.FAULT, got assert parts.intact() @@ -587,10 +708,10 @@ def test_a_failed_store_read_is_a_fault_not_a_miss(self) -> None: 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),))])) + (put,) = h.wait_for_results(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),))])) + (got,) = h.wait_for_results(h.get([RowIO(key, (reader.span(0),))])) assert got.status is RowStatus.FAULT, got assert reader.intact() @@ -607,9 +728,28 @@ def test_a_row_outside_every_registration_is_rejected_before_any_effect(self) -> 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())) + (result,) = h.wait_for_results(h.get(reader.rows())) h.assert_absent(result, reader, 0) + def test_a_batch_whose_extent_leaves_its_registration_is_rejected_before_any_effect( + self, + ) -> None: + """A batch's extents are checked, not only its rows: two rows each inside a registration + of its own give an extent inside neither, which the contract has the driver refuse. In + pool mode every endpoint lies in the one registration, so there is nothing to refuse.""" + if not self.registers_on_demand: + pytest.skip("a pool driver's endpoints share one registration") + with self.harness() as h: + h.needs(put=True, get=True, rows=2) # the two objects are published in one put + h.publish({"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + first, second = h.endpoint(["A"], fill=0x2D), h.endpoint(["B"], fill=0x2D) + with pytest.raises(DriverRejected): + h.get([RowIO(h.key("A"), (first.span(0),)), RowIO(h.key("B"), (second.span(0),))]) + assert first.read_row(0) == second.read_row(0) == b"-" * SLOT + for reader in (first, second): # each row is taken in a batch of its own + (result,) = h.wait_for_results(h.get(reader.rows())) + assert result.status is RowStatus.READY + 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) @@ -625,9 +765,23 @@ def test_a_span_leaving_or_outliving_its_registration_is_rejected(self) -> None: 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())) + (result,) = h.wait_for_results(h.get(reader.rows())) h.assert_absent(result, reader, 0) + def test_every_span_of_a_row_is_checked_against_the_registrations(self) -> None: + with self.harness() as h: + h.needs(put=True, get=True, multi_span=True) + inside = h.endpoint(["A"], fill=0x2D) + inside.fill_row(0, b"aaaaaaaa") + stray = HostBuffer(SLOT, fill=0x2D) # the second span of the row, in no registration + row = RowIO(h.key("M"), (inside.span(0), stray.span())) + with pytest.raises(DriverRejected): + h.put([row]) + with pytest.raises(DriverRejected): + h.get([row]) + assert stray.read() == b"-" * SLOT and stray.canaries_intact() + assert inside.read_row(0) == b"aaaaaaaa" + def test_a_device_region_is_rejected_unless_device_memory_is_declared(self) -> None: with self.harness() as h: h.needs(registers=True) @@ -640,22 +794,6 @@ def test_a_device_region_is_rejected_unless_device_memory_is_declared(self) -> N 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) @@ -668,7 +806,7 @@ def test_rows_beyond_the_declared_limit_are_rejected(self) -> None: with pytest.raises(DriverRejected): h.put(writer.rows()) reader = h.endpoint(["R0"], fill=0x2D) - (result,) = h.settle(h.get(reader.rows())) + (result,) = h.wait_for_results(h.get(reader.rows())) h.assert_absent(result, reader, 0) def test_an_undeclared_direction_is_rejected(self) -> None: @@ -695,15 +833,61 @@ def test_results_are_stable_and_access_ends_with_evidence(self) -> None: h.publish({"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) reader = h.endpoint(["A", "B"]) op = h.get(reader.rows()) - results = h.settle(op) + results = h.wait_for_results(op) assert [r.status for r in results] == [RowStatus.READY, RowStatus.READY] assert h.driver.poll(op) == results - h.settle_access(op) + h.wait_for_access_end(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: + def test_on_change_is_called_after_each_operation_settles(self) -> None: + """A caller that looks at an operation only when its ``on_change`` is called sees it + settle: some call comes after the last change shows.""" + with self.harness() as h: + h.needs(put=True, get=True, rows=2) + h.publish({"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + writer = h.endpoint(["C"]) + writer.fill_row(0, b"cccccccc") + reader = h.endpoint(["A", "B"], fill=0x2D) + for start, rows in ((h.put, writer.rows()), (h.get, reader.rows())): + seen: list[tuple[tuple[RowResult, ...], bool]] = [] + + def look(op: OpHandle, seen: list = seen) -> None: + seen.append((tuple(h.driver.poll(op)), h.driver.access_ended(op))) + + op = start(rows, on_change=look) + settled = (tuple(h.wait_for_results(op)), True) + h.wait_for_access_end(op) + wait_until(lambda: settled in seen, "on_change after the last change", h.timeout) + + def test_on_change_is_called_once_every_row_concluded_while_access_stays_open(self) -> None: + """A caller lands an outcome from the rows before their access ends, so the rows' settling + is told on its own.""" + staged = self.make_results_before_access_end() + if staged is None: + pytest.skip("this driver cannot conclude rows while access stays open") + driver, conclude, end_access = staged + with self.harness(driver) as h: + h.needs(get=True) + reader = h.endpoint(["R"], fill=0x2D) + told: list[tuple[bool, bool]] = [] + op = h.get( + reader.rows(), + on_change=lambda op: told.append( + (concluded(driver.poll(op)), driver.access_ended(op)) + ), + ) + conclude(op) + concluded_open = (True, False) + wait_until( + lambda: concluded_open in told, "on_change once every row concluded", h.timeout + ) + end_access(op) + + def test_an_unprovable_get_calls_on_change_keeps_its_region_and_drain_returns_at_once( + self, + ) -> None: staged = self.make_unprovable_get() if staged is None: pytest.skip("this driver cannot stage an operation whose end it cannot prove") @@ -711,12 +895,14 @@ def test_an_unprovable_get_makes_drain_return_at_once_and_keeps_its_region(self) 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 + told_unprovable: list[bool] = [] + op = h.get( + reader.rows(), + on_change=lambda op: told_unprovable.append(driver.access_unprovable(op)), ) + lose_the_evidence(op) + h.wait_for_results(op) + wait_until(lambda: True in told_unprovable, "on_change once unprovable", h.timeout) assert not driver.access_ended(op) started = time.monotonic() assert driver.drain(h.timeout) is False @@ -727,24 +913,84 @@ def test_an_unprovable_get_makes_drain_return_at_once_and_keeps_its_region(self) with pytest.raises(DriverRejected): reader.unregister() + def test_a_cancel_that_ends_an_operation_calls_its_on_change(self) -> None: + self._end_an_unstarted_get("cancel", lambda driver, op: driver.cancel(op)) + + def test_a_drain_that_ends_an_operation_calls_its_on_change(self) -> None: + self._end_an_unstarted_get("drain", lambda driver, op: driver.drain(0.0)) + + def _end_an_unstarted_get( + self, how: str, end: Callable[[StoreDriver, OpHandle], object] + ) -> None: + driver = self.make_driver_leaving_operations_unstarted() + if driver is None: + pytest.skip("this driver cannot hold an operation unstarted") + with self.harness(driver) as h: + h.needs(get=True, cancel=how == "cancel") + reader = h.endpoint(["Q"], fill=0x2D) + told: list[tuple[list[RowStatus], bool]] = [] + op = h.get( + reader.rows(), + on_change=lambda op: told.append((statuses(driver, op), driver.access_ended(op))), + ) + end(driver, op) + assert statuses(driver, op) == [RowStatus.CANCELLED] + ended = ([RowStatus.CANCELLED], True) + wait_until(lambda: ended in told, f"on_change once the {how} ended the get", h.timeout) + 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) + results = h.wait_for_results(op) h.driver.cancel(op) assert h.driver.poll(op) == results assert reader.read_row(0) == b"aaaaaaaa" + def test_a_put_returns_its_handle_before_the_store_answers(self) -> None: + """put runs on the caller's thread, the executor loop's: it hands the handle back while + the store still holds the row, and never waits for the store.""" + if not self.returns_before_the_store_answers: + pytest.skip("this driver answers inside put and get") + staged = self.make_blocked_put() + if staged is None: + pytest.skip("this driver cannot stage a put blocked in the store") + driver, blocked_key, blocked, unblock = staged + with self.harness(driver) as h: + h.needs(put=True) + source = h.endpoint(["P"]) + source.fill_row(0, b"pppppppp") + handed_back = threading.Event() + seen: dict[str, OpHandle] = {} + + def put() -> None: + seen["op"] = h.put([RowIO(blocked_key, (source.span(0),))]) + handed_back.set() + + putter = threading.Thread(target=put, daemon=True) + putter.start() + try: + assert handed_back.wait(h.timeout), "put waited for the store to answer" + if blocked is not None: + assert blocked.wait(h.timeout), "the put never reached the store" + op = seen["op"] + assert statuses(h.driver, op) == [RowStatus.PENDING] # still inside the store + assert h.driver.access_ended(op) is False + finally: + unblock() + putter.join(h.timeout) + (result,) = h.wait_for_results(seen["op"]) + assert result.status is RowStatus.READY + 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: + if not self.returns_before_the_store_answers: pytest.skip("this driver returns the handle of a get only once the store answered") staged = self.make_blocked_get() if staged is None: @@ -796,7 +1042,7 @@ def cancel() -> None: 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) + (after,) = h.wait_for_results(op) if during.concluded: assert after == during @@ -813,15 +1059,26 @@ def test_a_raise_after_the_operation_is_recorded_becomes_a_row_fault(self) -> No 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 + # Each declared direction is checked on its own: a get of an unpublished key still + # reaches the store, so a get-only driver is checked through its get alone. + capabilities = h.driver.capabilities + if not (capabilities.put or capabilities.get): + pytest.skip("this driver declares neither put nor get") + cases = [] + if capabilities.put: + writer = h.endpoint(["A"]) + writer.fill_row(0, b"aaaaaaaa") + cases.append(("put", h.put, writer.rows())) + if capabilities.get: + reader = h.endpoint(["A"], fill=0x2D) + cases.append(("get", h.get, reader.rows())) + for direction, start, rows in cases: + try: + op = start(rows) + except Exception as error: # noqa: BLE001 - the check is that nothing escapes + pytest.fail(f"{direction} raised after the operation was recorded: {error!r}") + (result,) = h.wait_for_results(op) + assert result.status is RowStatus.FAULT, (direction, 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() @@ -835,20 +1092,20 @@ def test_a_key_held_by_an_unfinished_write_is_not_ready_for_get_put_and_probe(se writer.fill_row(0, b"hhhhhhhh") write = [RowIO(key, (writer.span(0),))] try: - (got,) = h.settle(h.get([RowIO(key, (early.span(0),))])) + (got,) = h.wait_for_results(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)) + (put,) = h.wait_for_results(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)) + (put,) = h.wait_for_results(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),))])) + (got,) = h.wait_for_results(h.get([RowIO(key, (later.span(0),))])) assert got.status is RowStatus.READY assert later.read_row(0) == b"hhhhhhhh" @@ -862,20 +1119,24 @@ def test_a_row_with_no_room_is_declined_while_the_others_go_on(self) -> None: writer.fill_row(0, b"aaaaaaaa") writer.fill_row(1, b"bbbbbbbb") op = h.put(writer.rows()) - results = h.settle(op) + results = h.wait_for_results(op) assert sorted((r.status for r in results), key=lambda s: s.value) == [ RowStatus.DECLINED, RowStatus.READY, ], results - h.settle_access(op) + h.wait_for_access_end(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())) + got = h.wait_for_results(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: + def test_a_get_blocked_in_the_store_delays_nothing_and_keeps_its_memory(self) -> None: + """A get blocked inside the store delays no other call, nor its own handle when the + driver returns handles before the store answers; until the store lets it go, its access + stays open, through a drain too, and its region refuses to unregister (for a driver that + registers on demand).""" staged = self.make_blocked_get() if staged is None: pytest.skip("this driver cannot stage a get blocked in the store") @@ -900,12 +1161,12 @@ def slow_get() -> None: def quick_calls() -> None: try: op = h.get(quick.rows()) - seen["quick"] = h.settle(op) + seen["quick"] = h.wait_for_results(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: + if self.returns_before_the_store_answers: assert handle_back.wait(h.timeout), "the slow get kept its handle" if handle_back.is_set(): slow_op = seen["slow_op"] @@ -921,6 +1182,8 @@ def quick_calls() -> None: else: seen["unregister"] = "accepted" seen["drain"] = h.driver.drain(0.05) + if handle_back.is_set(): + seen["slow_ended_after_drain"] = h.driver.access_ended(seen["slow_op"]) except BaseException as error: # noqa: BLE001 - reported from the main thread failures.append(error) @@ -947,10 +1210,11 @@ def quick_calls() -> None: 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: + if self.returns_before_the_store_answers: 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 + assert seen["slow_ended_after_drain"] is False, "drain forced an access end" slow_results = seen["slow_poll"] assert isinstance(slow_results, tuple) assert all(r.status is not RowStatus.READY for r in slow_results) @@ -1023,17 +1287,25 @@ def test_probe_raises_on_a_store_error_rather_than_answering(self) -> None: # 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.""" + """Also for a driver whose memory was registered up front: its one registration, or the + batch it was registered as, is revoked twice through the call that made it, 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) + revoke = h.driver.unregister + if self.registers_together and not self.registers_on_demand: + assert isinstance(h.driver, RegistersRegionsTogether) + revoke = h.driver.unregister_regions + revoke(endpoint.handle) + try: + revoke(endpoint.handle) + except Exception as error: # noqa: BLE001 - the check is that nothing escapes + pytest.fail(f"a second unregister raised {error!r}") endpoint.released = True # the harness must not give pool memory back twice - def test_unregister_frees_the_range_for_a_new_registration(self) -> None: + def test_a_registered_range_cannot_be_registered_again_until_unregistered(self) -> None: with self.harness() as h: h.needs(registers=True) endpoint = h.endpoint(["A"]) @@ -1046,6 +1318,145 @@ def test_unregister_frees_the_range_for_a_new_registration(self) -> None: pytest.fail(f"the range stayed registered: {error}") h.driver.unregister(handle) + def test_regions_registered_together_are_revoked_together_and_the_range_freed(self) -> None: + """Through the optional ``RegistersRegionsTogether``: every region of a batch is live once + registered, each overlaps a new registration until the batch is revoked, revocation is + idempotent, and the whole range is free after.""" + staged = self.make_together_registration() + if staged is None: + pytest.skip("this suite stages no batch registration") + driver, memory = staged + assert isinstance(driver, RegistersRegionsTogether) + assert len(memory) >= 2, "a batch check needs at least two regions" + specs = specs_of(memory) + released = False + try: + handle = driver.register_regions(specs) + for index, spec in enumerate(specs): + try: + stray = driver.register_regions([spec]) + except DriverRejected: + continue # the region is live: it overlaps a new registration + driver.unregister_regions(stray) + pytest.fail(f"region {index} of the batch was not registered") + for index, region in enumerate(memory): # the registrations are real: rows move + if not _prove_registered(driver, region, index, self.timeout): + _skip_unprovable(index) + driver.unregister_regions(handle) + _revoke_again(driver, handle) + try: + again = driver.register_regions(specs) + except DriverRejected as error: + pytest.fail(f"the revoked batch kept its range: {error}") + driver.unregister_regions(again) + assert driver.drain(self.timeout) is True + driver.close() + released = True + finally: + if not released: + UNRELEASED_MEMORY.extend(memory) + + def test_a_batch_is_not_revoked_while_any_of_its_regions_is_busy(self) -> None: + """``unregister_regions`` refuses while an operation may still touch any region of the + batch and leaves every region registered, the idle ones included; once access ended it + succeeds, idempotently.""" + staged = self.make_busy_together_registration() + if staged is None: + pytest.skip("this suite stages no batch with a stalled get") + driver, memory, key, blocked, unblock = staged + assert isinstance(driver, RegistersRegionsTogether) + assert len(memory) >= 2, "a batch check needs at least two regions" + if not driver.capabilities.get: + pytest.skip("this driver declares no get") + specs = specs_of(memory) + slow: dict[str, OpHandle] = {} + row = RowIO(key, (Span(memory[-1].address, SLOT),)) # the last region is the busy one + getter = threading.Thread( + target=lambda: slow.update(op=driver.get(RowBatch.of([row]))), daemon=True + ) + refusal: list[BaseException | None] = [] + + def revoke_while_busy() -> None: + try: + driver.unregister_regions(handle) + except BaseException as error: # noqa: BLE001 - judged on the main thread + refusal.append(error) + else: + refusal.append(None) + + # Off this thread: a revocation that waits for the transfer would wait for the gate this + # thread holds, and must be a failure rather than a hang. + revoker = threading.Thread(target=revoke_while_busy, daemon=True) + released = False + try: + handle = driver.register_regions(specs) + getter.start() # a blocking driver's get returns only once the store lets it + assert blocked.wait(self.timeout), "the get never reached the store" + revoker.start() + revoker.join(self.timeout) + assert not revoker.is_alive(), "unregister_regions waited for the transfer" + assert isinstance(refusal[0], DriverRejected), ( + f"a busy batch was revoked: {refusal[0]!r}" + ) + for index, region in enumerate(memory[:-1]): # the idle regions stayed registered + if not _prove_registered(driver, region, index, self.timeout): + _skip_unprovable(index) + unblock() + getter.join(self.timeout) + assert not getter.is_alive(), "the stalled get never returned" + if not _access_settled(driver, slow["op"], self.timeout): + _skip_unprovable(len(memory) - 1) + driver.unregister_regions(handle) + _revoke_again(driver, handle) + assert driver.drain(self.timeout) is True + driver.close() + released = True + finally: + unblock() + for worker in (getter, revoker): + if worker.ident is not None: # started: let it leave before the check ends + worker.join(self.timeout) + if not released: + UNRELEASED_MEMORY.extend(memory) + + def test_a_failed_batch_registration_is_revocable_or_left_nothing_behind(self) -> None: + """A ``register_regions`` that fails after the store was asked either certifies that + nothing of the batch is registered, so the same memory registers again at once, or + hands out a handle that revokes whatever is, idempotently. A failure of any other shape + is allowed and leaves the memory protected; nothing more is checked for it.""" + staged = self.make_failing_together_registrations() + if not staged: + pytest.skip("this suite stages no failing batch registration") + for driver, memory in staged: + assert isinstance(driver, RegistersRegionsTogether) + specs = specs_of(memory) + released = False + try: + try: + driver.register_regions(specs) + except DriverRejected: + handle = None + except RegistrationError as error: + handle = error.handle + except Exception: # noqa: BLE001 - some of the batch may be registered for good + continue + else: + pytest.fail("the staged registration did not fail") + if handle is not None: + driver.unregister_regions(handle) + _revoke_again(driver, handle) + try: + again = driver.register_regions(specs) + except DriverRejected as error: + pytest.fail(f"something of the failed batch stayed registered: {error}") + driver.unregister_regions(again) + assert driver.drain(self.timeout) is True + driver.close() + released = True + finally: + if not released: + UNRELEASED_MEMORY.extend(memory) + def test_drain_is_one_way(self) -> None: with self.harness() as h: h.needs(put=True, get=True) @@ -1078,3 +1489,8 @@ def test_close_needs_a_true_drain_and_no_live_registrations(self) -> None: h.done = True # closed by the check itself with pytest.raises(DriverRejected): h.put(writer.rows()) + + +def _batch(rows: Sequence[RowIO]) -> RowBatch: + """``rows`` as the batch a driver takes; a check that builds its rows by hand lists them.""" + return rows if isinstance(rows, RowBatch) else RowBatch.of(rows) diff --git a/tensorrt_llm/_torch/disaggregation/testing/memory_store.py b/tensorrt_llm/_torch/disaggregation/testing/memory_store.py index 88471bdad7df..6daf2c486429 100644 --- a/tensorrt_llm/_torch/disaggregation/testing/memory_store.py +++ b/tensorrt_llm/_torch/disaggregation/testing/memory_store.py @@ -9,14 +9,21 @@ 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. +In automatic mode ``put``/``get`` run the operation before returning; in background mode they +start a thread of the driver's own to run it. Neither keeps the protocol's rule that ``put`` and +``get`` only record and queue the work, which a test can afford, and automatic mode waits on the +store, so over a store a test holds, use background mode. 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. + +``TogetherMemoryDriver`` is the same driver registering regions in batches, for the optional +``RegistersRegionsTogether`` protocol; its single-region ``register`` and ``unregister`` are +rejected, as for a store that opens its client with the memory it will use. """ from __future__ import annotations @@ -26,18 +33,21 @@ import threading from collections.abc import Iterable, Sequence from dataclasses import dataclass -from typing import Literal +from typing import Any, Literal, NoReturn from tensorrt_llm._torch.disaggregation.base.memory import MemoryKind from tensorrt_llm._torch.disaggregation.blob.store import ( DriverCapabilities, DriverRejected, Hint, + OnChange, OpHandle, ProbeStatus, RegionHandle, RegionLayout, + RegionSpec, RegistrationError, + RowBatch, RowIO, RowResult, RowStatus, @@ -205,6 +215,7 @@ class MemoryOp: kind: Kind rows: list[_Row] hint: Hint = None + on_change: OnChange | None = None access_ended: bool = False unprovable: bool = False @@ -233,7 +244,7 @@ def _scatter(spans: Sequence[Span], payload: bytes) -> None: class MemoryDriver: - """A ``StoreDriver`` over a ``MemoryStore``; see the module docstring for the two modes.""" + """A ``StoreDriver`` over a ``MemoryStore``; see the module docstring for its three modes.""" def __init__( self, @@ -270,19 +281,8 @@ 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) + registration = _Registration(self._validate_region(address, nbytes, layout, ()), layout) if self._register_failure is not None: reason, escaped = self._register_failure self._register_failure = None @@ -305,11 +305,13 @@ def unregister(self, handle: RegionHandle) -> None: raise StoreError(reason) registration.live = False - def put(self, rows: Sequence[RowIO]) -> MemoryOp: - return self._submit("put", rows) + def put(self, rows: RowBatch, *, on_change: OnChange | None = None) -> MemoryOp: + return self._submit("put", rows, on_change=on_change) - def get(self, rows: Sequence[RowIO], hint: Hint = None) -> MemoryOp: - return self._submit("get", rows, hint) + def get( + self, rows: RowBatch, hint: Hint = None, *, on_change: OnChange | None = None + ) -> MemoryOp: + return self._submit("get", rows, hint, on_change) def poll(self, op: OpHandle) -> tuple[RowResult, ...]: with self._lock: @@ -324,9 +326,15 @@ def access_unprovable(self, op: OpHandle) -> bool: return self._op(op).unprovable def cancel(self, op: OpHandle) -> None: + """Cancels the rows not started yet, and tells ``on_change`` only when that changed one.""" with self._lock: + op = self._op(op) + before = ([row.published for row in op.rows], op.access_ended) if self._cancel_effective: - self._cancel_unstarted(self._op(op)) + self._cancel_unstarted(op) + changed = ([row.published for row in op.rows], op.access_ended) != before + if changed: + self._call_on_change(op) def probe( self, keys: Sequence[str], budget_s: float, hint: Hint = None @@ -371,8 +379,12 @@ def drain(self, budget_s: float) -> bool: with self._lock: self._require_not_closed() self._start_draining() - for op in self._ops: + ops_to_cancel = [op for op in self._ops if any(_cancellable(row) for row in op.rows)] + for op in ops_to_cancel: self._cancel_unstarted(op) + for op in ops_to_cancel: # the operations this drain changed; outside the lock + self._call_on_change(op) + with self._lock: self._lock.wait_for(self._drain_can_return, timeout=budget_s) return self._quiet() @@ -420,6 +432,7 @@ def publish_results(self, op: MemoryOp, rows: Iterable[int] | None = None) -> No if row.staged is None: raise RuntimeError(f"row {index} has no staged result to publish") row.published, row.staged = row.staged, None + self._call_on_change(op) def complete(self, op: MemoryOp, rows: Iterable[int] | None = None) -> None: """Transfer ``rows`` and publish their results, leaving access open.""" @@ -437,6 +450,7 @@ def fail(self, op: MemoryOp, reason: str, rows: Iterable[int] | None = None) -> if row.started or row.published.concluded: raise RuntimeError(f"row {index} was already taken or decided") row.published = faulted(reason) + self._call_on_change(op) 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 @@ -460,6 +474,7 @@ def end_access(self, op: MemoryOp) -> None: raise RuntimeError("a transfer is still in progress") op.access_ended = True self._lock.notify_all() + self._call_on_change(op) def mark_unprovable(self, op: MemoryOp) -> None: """Declare that the driver will never learn when this operation's access ends, as after a @@ -468,6 +483,7 @@ def mark_unprovable(self, op: MemoryOp) -> None: self._require_open(op) op.unprovable = True self._lock.notify_all() + self._call_on_change(op) def set_cancel_effective(self, effective: bool) -> None: self._cancel_effective = effective @@ -518,17 +534,10 @@ def _check_hint(self, hint: Hint) -> None: 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") @@ -548,12 +557,40 @@ 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 _validate_region( + self, address: int, nbytes: int, layout: RegionLayout, taken: Sequence[Span] + ) -> Span: + """The span this driver would register for the region, or DriverRejected. ``taken`` are + the spans of the same batch, which count as live registrations.""" + 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) + 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}") + for other_span in taken: + if other_span.overlaps(span): + raise DriverRejected(f"region {span} overlaps another of the batch {other_span}") + return span + 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_extents(self, rows: Sequence[RowIO]) -> None: + """A ``RowBatch`` promises each extent inside one registered region, which a driver may + check in place of every row; rows given as a plain sequence carry no extent.""" + for extent in getattr(rows, "extents", ()): + if not self._registered(extent): + raise DriverRejected(f"extent {extent} is outside every live registration") + def _check_region_idle(self, registration: _Registration) -> None: busy = sum( 1 @@ -595,18 +632,24 @@ def _quiet(self) -> bool: # Internals. - def _submit(self, kind: Kind, rows: Sequence[RowIO], hint: Hint = None) -> MemoryOp: + def _submit( + self, + kind: Kind, + rows: Sequence[RowIO], + hint: Hint = None, + on_change: OnChange | None = 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._check_extents(rows) + op = MemoryOp(kind, [_Row(io) for io in rows], hint, on_change) self._ops.append(op) if hint is not None: self.hints_seen.append(hint) @@ -627,12 +670,17 @@ def _submit(self, kind: Kind, rows: Sequence[RowIO], hint: Hint = None) -> Memor self._workers.append(_Worker(thread, finished)) if not self.background and not self.controlled: self._run(op, None) + elif op.access_ended: # no worker thread: it concluded at once + self._call_on_change(op) return op def _run(self, op: MemoryOp, finished: threading.Event | None) -> None: try: self._before_execute(op) self._execute(op) + # Told before the worker counts as finished: a drain that returned True earlier would + # let the owner close the driver while the callback still queries it. + self._call_on_change(op) finally: if finished is not None: finished.set() @@ -703,6 +751,11 @@ 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 _call_on_change(self, op: MemoryOp) -> None: + """Outside the lock, as ``OnChange`` asks.""" + if op.on_change is not None: + op.on_change(op) + 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: @@ -712,7 +765,7 @@ def _fail_unstarted(self, op: MemoryOp, reason: str) -> None: def _cancel_unstarted(self, op: MemoryOp) -> None: for row in op.rows: - if not row.started and not row.published.concluded: + if _cancellable(row): row.published = _CANCELLED self._settle_untouched(op) @@ -755,3 +808,95 @@ def _require_open(self, op: MemoryOp) -> None: @staticmethod def _indexes(op: MemoryOp, rows: Iterable[int] | None) -> list[int]: return list(range(len(op.rows))) if rows is None else list(rows) + + +def _cancellable(row: _Row) -> bool: + """Neither taken by a transfer nor decided: a cancel or a drain may still cancel it.""" + return not row.started and not row.published.concluded + + +@dataclass(eq=False) +class _Batch: + """The handle of one ``register_regions`` call: its registrations, revoked together.""" + + registrations: list[_Registration] + live: bool = True + + +class TogetherMemoryDriver(MemoryDriver): + """The in-memory driver as a ``RegistersRegionsTogether``: regions come and go in batches, + one handle per ``register_regions`` call, and the single-region ``register`` and + ``unregister`` are rejected. Everything else is the in-memory driver's, controls included.""" + + def __init__(self, store: MemoryStore, **options: Any) -> None: + super().__init__(store, **options) + self._batches: list[_Batch] = [] + self._batch_failure: tuple[str, bool] | None = None + self._revocation_failure: str | None = None + + def register(self, address: int, nbytes: int, layout: RegionLayout) -> NoReturn: + raise DriverRejected("this driver registers regions together; use register_regions") + + def unregister(self, handle: RegionHandle) -> NoReturn: + raise DriverRejected("this driver revokes regions together; use unregister_regions") + + def register_regions(self, specs: Sequence[RegionSpec]) -> _Batch: + if not specs: + raise DriverRejected("a batch registers at least one region") + with self._lock: + spans: list[Span] = [] + for spec in specs: # all checked before any is kept: a rejection escapes nothing + spans.append(self._validate_region(spec.address, spec.nbytes, spec.layout, spans)) + batch = _Batch( + [_Registration(span, spec.layout) for span, spec in zip(spans, specs, strict=True)] + ) + if self._batch_failure is not None: + reason, escaped = self._batch_failure + self._batch_failure = None + if not escaped: + raise RegistrationError(reason) + self._keep(batch) # the store kept the batch before it answered + raise RegistrationError(reason, batch) + self._keep(batch) + return batch + + def unregister_regions(self, handle: RegionHandle) -> None: + with self._lock: + batch = self._batch(handle) + if not batch.live: + return + for registration in batch.registrations: + self._check_region_idle(registration) + if self._revocation_failure is not None: + reason, self._revocation_failure = self._revocation_failure, None + raise StoreError(reason) + for registration in batch.registrations: + registration.live = False + batch.live = False + + # Controls. + + def fail_next_register_regions(self, reason: str, *, escaped: bool) -> None: + """Make the next ``register_regions`` raise; with ``escaped`` the store keeps the batch + anyway and the error carries its handle.""" + self._batch_failure = (reason, escaped) + + def fail_next_unregister_regions(self, reason: str) -> None: + """Make the next ``unregister_regions`` of a live batch raise, the batch staying live.""" + self._revocation_failure = reason + + def live_batches(self) -> list[_Batch]: + with self._lock: + return [batch for batch in self._batches if batch.live] + + # Internals. + + def _keep(self, batch: _Batch) -> None: + self._registrations.extend(batch.registrations) + self._batches.append(batch) + + def _batch(self, handle: RegionHandle) -> _Batch: + self._require_not_closed() + if not isinstance(handle, _Batch) or handle not in self._batches: + raise ValueError("handle does not belong to this driver") + return handle diff --git a/tests/unittest/_torch/disaggregation/blob/cpu_fixtures.py b/tests/unittest/_torch/disaggregation/blob/cpu_fixtures.py new file mode 100644 index 000000000000..f130e5dd8396 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/cpu_fixtures.py @@ -0,0 +1,589 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU stand-ins for the lender side of a blob backend, and ways for a test to steer a driver. + +``Endpoint`` is one side of a transfer: host slots in one registered region, a driver, a region +binding, a backend and the waits a test lease holder needs. ``FakeLeaseHolder`` is that holder: +one decision per attempt, the lease back on quiescence. ``KeyGatedStore`` stops store calls at +its gate until the test opens it. ``DriverSpy`` forwards to a real driver, records every call +with its thread, and can gate a method on an event or make it raise. Every wait has an upper bound. +""" + +from __future__ import annotations + +import hashlib +import threading +import time +from collections.abc import Callable, Iterable, Sequence +from dataclasses import dataclass + +import numpy as np + +from tensorrt_llm._torch.disaggregation.base.backend import MemoryKind, Outcome +from tensorrt_llm._torch.disaggregation.blob.backend import ( + BlobAttempt, + BlobBackend, + BlobConfig, + RegionBinding, +) +from tensorrt_llm._torch.disaggregation.blob.keys import encode_key +from tensorrt_llm._torch.disaggregation.blob.store import ( + DriverCapabilities, + Hint, + OnChange, + OpHandle, + RegionHandle, + RegionLayout, + RegionSpec, + RegistersRegionsTogether, + RegistrationError, + RowBatch, + RowIO, + RowResult, + StoreDriver, +) +from tensorrt_llm._torch.disaggregation.testing.host_memory import HostBuffer +from tensorrt_llm._torch.disaggregation.testing.memory_store import ( + Commit, + MemoryDriver, + MemoryStore, + StoredObject, +) +from tensorrt_llm._torch.disaggregation.testing.object_trailer import Checksum, encode_trailer + +WAIT_S = 5.0 +"""Upper bound of every wait in these tests; reached only when the code under test is wrong.""" + +GATED_CALL_S = 10 * WAIT_S +"""How long a gated call waits for its gate: past any test wait, so a gate left shut fails the +test's own wait first, yet a failed test never strands the thread.""" + +PROMPT_S = 0.5 +"""A call that must not wait on the driver returns well within this.""" + +KEY_PREFIX = "t" +SLOT_BYTES = 8 +OBJECT_KEY_WIDTH = 16 + + +def object_key(name: str) -> np.ndarray: + """The lender-style opaque object key of a test row name.""" + digest = hashlib.sha256(name.encode()).digest()[:OBJECT_KEY_WIDTH] + return np.frombuffer(digest, dtype=np.uint8) + + +def store_key(name: str) -> str: + """The store key the backend derives for a test row name.""" + return encode_key(KEY_PREFIX, object_key(name).tobytes()) + + +def commit_object(store: MemoryStore, name: str, data: bytes) -> Commit: + """Commits ``data`` straight into ``store`` under the key the backend derives for ``name``, + with the trailer ``MemoryDriver`` keeps beside each object.""" + key = store_key(name) + trailer = encode_trailer(key.encode(), len(data), checksum=Checksum.CRC32, payload=data) + return store.commit(key, data, trailer) + + +def masks(outcome: Outcome) -> list[list[bool]]: + return [mask.tolist() for mask in outcome.delivered_masks] + + +def wait_until(predicate: Callable[[], bool], what: str, timeout: float = WAIT_S) -> None: + """Checks ``predicate`` every millisecond until it holds; fails the test after ``timeout``.""" + deadline = time.monotonic() + timeout + while not predicate(): + if time.monotonic() > deadline: + raise AssertionError(f"timed out waiting for {what}") + time.sleep(0.001) + + +class _Sealable: + """Raises on any attribute read once sealed, to prove nothing reads it any more.""" + + _sealed = False + + def seal(self) -> None: + object.__setattr__(self, "_sealed", True) + + def __getattribute__(self, name: str): + if name != "_sealed" and object.__getattribute__(self, "_sealed"): + raise AssertionError(f"{type(self).__name__}.{name} read after the view was sealed") + return object.__getattribute__(self, name) + + +class CpuRegion: + """``slot_count`` slots of ``slot_bytes`` in host memory: a ``SlotRegion``. ``memory_kind`` is + what the region declares; DEVICE stands for device memory in-process.""" + + def __init__( + self, + slot_bytes: int, + slot_count: int, + fill: int = 0, + layout_label: str = "kv", + memory_kind: MemoryKind = MemoryKind.HOST, + ) -> None: + self.buffer = HostBuffer(slot_bytes * max(slot_count, 1), fill) + self.layout_label = layout_label + self.memory_kind = memory_kind + self.address = self.buffer.address + self.nbytes = self.buffer.nbytes + self.slot_bytes = slot_bytes + self.slot_count = slot_count + + def slot(self, index: int) -> int: + return self.address + index * self.slot_bytes + + def slots(self, *indexes: int) -> np.ndarray: + """The addresses of slots ``indexes``, as a view's ``addresses``.""" + return np.array([self.slot(index) for index in indexes], dtype=np.int64) + + def read(self, index: int) -> bytes: + return self.buffer.read(index * self.slot_bytes, self.slot_bytes) + + def write(self, index: int, data: bytes) -> None: + if len(data) != self.slot_bytes: + raise ValueError(f"a slot holds {self.slot_bytes} bytes, got {len(data)}") + self.buffer.write(data, index * self.slot_bytes) + + +class Rows(_Sealable): + """One group of a view: a ``LayerGroupRows``. ``rows`` overrides ``len`` for malformed views.""" + + def __init__( + self, + layer_group_index: int, + object_keys: np.ndarray | None, + addresses: np.ndarray | None, + region_index: int | None, + rows: int | None = None, + ) -> None: + self.layer_group_index = layer_group_index + self.object_keys = object_keys + self.addresses = addresses + self.region_index = region_index + self.rows = rows + + def __len__(self) -> int: + if self.rows is not None: + return self.rows + return 0 if self.addresses is None else len(self.addresses) + + +class CpuView(_Sealable): + """A ``View`` over ``Rows`` groups; ``seal`` makes any later read of it or its groups raise.""" + + def __init__(self, groups: Iterable[Rows]) -> None: + self.groups = tuple(groups) + + def row_masks(self, fill_value: bool = False) -> tuple[np.ndarray, ...]: + return tuple(np.full(len(group), bool(fill_value)) for group in self.groups) + + def seal(self) -> None: + for group in self.groups: + group.seal() + super().seal() + + +class FakeHold: + """A lender region hold that records each release and its thread.""" + + def __init__(self) -> None: + self.release_threads: list[int] = [] + + def release(self) -> None: + self.release_threads.append(threading.get_ident()) + + @property + def released(self) -> bool: + return bool(self.release_threads) + + +class WakeCounter: + """A ``WakeFn`` that counts calls; tests wait on it for state the backend published.""" + + def __init__(self, on_wake: Callable[[], None] | None = None) -> None: + self.count = 0 + self.on_wake = on_wake + self._woken = threading.Condition() + + def __call__(self) -> None: + if self.on_wake is not None: + self.on_wake() + with self._woken: + self.count += 1 + self._woken.notify_all() + + def wait_for(self, predicate: Callable[[], bool], what: str) -> None: + """Waits for ``predicate``, re-checked on every wake. State that lands without a wake times + the test out only when it lands after the wait began; a test that pins a wake counts + the wakes itself.""" + with self._woken: + assert self._woken.wait_for(predicate, WAIT_S), f"no wake brought {what}" + + +class KeyGatedStore(MemoryStore): + """Stops ``get``s and/or ``commit``s of the chosen keys (all by default) at ``gate`` until the + test sets it; ``calls_at_gate`` counts the calls that reached it.""" + + def __init__( + self, *, gets: bool = False, commits: bool = False, keys: Iterable[str] | None = None + ) -> None: + super().__init__() + self.gate = threading.Event() + self.calls_at_gate = 0 + self._gets = gets + self._commits = commits + self._keys = None if keys is None else set(keys) + self._arrived = threading.Condition() + + def get(self, key: str) -> StoredObject | None: + if self._gets: + self._wait_at_gate(key) + return super().get(key) + + def commit(self, key: str, payload: bytes, trailer: bytes) -> Commit: + if self._commits: + self._wait_at_gate(key) + return super().commit(key, payload, trailer) + + def wait_for_calls_at_gate(self, count: int) -> None: + with self._arrived: + arrived = self._arrived.wait_for(lambda: self.calls_at_gate >= count, WAIT_S) + assert arrived, f"{self.calls_at_gate} of {count} store calls reached the gate" + + def _wait_at_gate(self, key: str) -> None: + if self.gate.is_set() or (self._keys is not None and key not in self._keys): + return + with self._arrived: + self.calls_at_gate += 1 + self._arrived.notify_all() + self.gate.wait(GATED_CALL_S) + + +class RollbackFailsDriver(MemoryDriver): + """Registers the first region with a revocation that will fail, then refuses the second.""" + + def register(self, address: int, nbytes: int, layout: RegionLayout) -> RegionHandle: + if self.live_registrations(): + raise RegistrationError("out of memory regions") + handle = super().register(address, nbytes, layout) + self.fail_next_unregister(handle, "device busy") + return handle + + +@dataclass(frozen=True) +class Call: + """One call into a ``DriverSpy``.""" + + method: str + thread: str + rows: tuple[RowIO, ...] = () + op: object = None + + +class DriverSpy: + """A ``StoreDriver`` that forwards to ``inner`` and records every call with its thread name. + + ``gate(method, after)`` makes calls of ``method`` past the first ``after`` wait until the + returned event is set; ``fail(method, error, times)`` makes the next ``times`` calls raise + ``error`` instead; + ``answer(method, value, times)`` makes them return ``value`` without reaching the driver. + ``ops`` lists every operation ``put``/``get`` returned, with its rows; ``call_on_change(op)`` + calls the ``on_change`` that came with it again, as a driver may. + """ + + def __init__(self, inner: StoreDriver) -> None: + self.inner = inner + self.calls: list[Call] = [] + self.ops: list[tuple[tuple[RowIO, ...], OpHandle]] = [] + self._on_changes: dict[OpHandle, OnChange | None] = {} + self._changed = threading.Condition() + self._gates: dict[str, tuple[threading.Event, int]] = {} # method -> (gate, after) + self._failures: dict[str, list] = {} + self._answers: dict[str, list] = {} + + def gate(self, method: str, after: int = 0) -> threading.Event: + """Makes every call of ``method`` past the first ``after`` wait on the returned event.""" + gate = threading.Event() + self._gates[method] = (gate, after) + return gate + + def fail(self, method: str, error: Exception, times: int = 1) -> None: + self._failures[method] = [error, times] + + def answer(self, method: str, value: object, times: int = 1) -> None: + self._answers[method] = [value, times] + + def op_with(self, key: str) -> OpHandle: + """The latest operation carrying the row stored under ``key``, once ``put`` or ``get`` + returned it.""" + found: list[OpHandle] = [] + + def returned() -> bool: + found[:] = [op for rows, op in self.ops if any(row.key == key for row in rows)] + return bool(found) + + self.wait_for(returned, f"an operation for {key}") + return found[-1] + + def call_on_change(self, op: OpHandle) -> None: + """Calls ``op``'s ``on_change`` on this thread, as a driver may.""" + with self._changed: + on_change = self._on_changes[op] + on_change(op) + + def calls_to(self, method: str, op: object = None) -> list[Call]: + with self._changed: + return [c for c in self.calls if c.method == method and (op is None or c.op is op)] + + def wait_for(self, predicate: Callable[[], bool], what: str) -> None: + with self._changed: + assert self._changed.wait_for(predicate, WAIT_S), f"the driver never saw {what}" + + @property + def capabilities(self) -> DriverCapabilities: + return self.inner.capabilities + + def register(self, address: int, nbytes: int, layout: RegionLayout) -> RegionHandle: + self._enter("register") + return self.inner.register(address, nbytes, layout) + + def unregister(self, handle: RegionHandle) -> None: + self._enter("unregister") + self.inner.unregister(handle) + + def put(self, rows: RowBatch, *, on_change: OnChange | None = None) -> OpHandle: + self._enter("put", rows=rows) + return self._record_op(rows, self.inner.put(rows, on_change=on_change), on_change) + + def get( + self, rows: RowBatch, hint: Hint = None, *, on_change: OnChange | None = None + ) -> OpHandle: + self._enter("get", rows=rows) + if hint is None: + op = self.inner.get(rows, on_change=on_change) + else: + op = self.inner.get(rows, hint, on_change=on_change) + return self._record_op(rows, op, on_change) + + def poll(self, op: OpHandle) -> tuple[RowResult, ...]: + self._enter("poll", op=op) + scripted = self._scripted_answer("poll") + return self.inner.poll(op) if scripted is None else scripted[0] + + def access_ended(self, op: OpHandle) -> bool: + self._enter("access_ended", op=op) + return self.inner.access_ended(op) + + def access_unprovable(self, op: OpHandle) -> bool: + self._enter("access_unprovable", op=op) + return self.inner.access_unprovable(op) + + def cancel(self, op: OpHandle) -> None: + self._enter("cancel", op=op) + self.inner.cancel(op) + + def drain(self, budget_s: float) -> bool: + self._enter("drain") + scripted = self._scripted_answer("drain") + return self.inner.drain(budget_s) if scripted is None else scripted[0] + + def close(self) -> None: + self._enter("close") + self.inner.close() + + def _scripted_answer(self, method: str) -> list | None: + with self._changed: + scripted = self._answers.get(method) + if scripted is None or scripted[1] <= 0: + return None + scripted[1] -= 1 + return scripted + + def _record_op( + self, rows: Sequence[RowIO], op: OpHandle, on_change: OnChange | None + ) -> OpHandle: + with self._changed: + self.ops.append((tuple(rows), op)) + self._on_changes[op] = on_change + self._changed.notify_all() + return op + + def _enter(self, method: str, rows: Sequence[RowIO] = (), op: object = None) -> None: + with self._changed: + self.calls.append(Call(method, threading.current_thread().name, tuple(rows), op)) + self._changed.notify_all() + failure = self._failures.get(method) + if failure is not None and failure[1] > 0: + failure[1] -= 1 + else: + failure = None + gate = self._gates.get(method) + entered = sum(1 for call in self.calls if call.method == method) # this one included + if gate is not None and entered > gate[1]: + gate[0].wait(GATED_CALL_S) + if failure is not None: + raise failure[0] + + +class TogetherDriverSpy(DriverSpy): + """A ``DriverSpy`` over a driver that registers regions together. The spy implements the + protocol too, so a binding over it takes the batch path; ``gate`` and ``fail`` work on + ``register_regions`` and ``unregister_regions`` as on any other method.""" + + def __init__(self, inner: StoreDriver) -> None: + assert isinstance(inner, RegistersRegionsTogether) + super().__init__(inner) + self.together = inner + + def register_regions(self, specs: Sequence[RegionSpec]) -> RegionHandle: + self._enter("register_regions") + return self.together.register_regions(specs) + + def unregister_regions(self, handle: RegionHandle) -> None: + self._enter("unregister_regions") + self.together.unregister_regions(handle) + + +class Endpoint: + """One side of a transfer: host slots in one registered region, a driver, a backend, waits. + + ``others`` are further regions of the lender, after the endpoint's own (region 0); the binding + registers those the driver can access. Slot ``i`` belongs to the row named ``names[i]``. + ``driver`` is a ``DriverSpy`` over the driver passed in: a test steers the one it passed + (``complete``, ``end_access``) and asks the spy what was called. ``close`` is the clean end + every test that finishes its transfers reaches: shutdown returns no labels, the hold is released + on this thread and every backend thread has exited. + """ + + def __init__( + self, + driver: StoreDriver, + names: Sequence[str], + *, + slot_bytes: int = SLOT_BYTES, + fill: int = 0, + on_wake: Callable[[], None] | None = None, + others: Sequence[CpuRegion] = (), + ) -> None: + self.names = list(names) + self.region = CpuRegion(slot_bytes, len(self.names), fill) + spy = TogetherDriverSpy if isinstance(driver, RegistersRegionsTogether) else DriverSpy + self.driver = spy(driver) + self.hold = FakeHold() + self.binding = RegionBinding((self.region, *others), self.driver, self.hold) + self.binding.register() + self.wake = WakeCounter(on_wake) + before = set(threading.enumerate()) + self.backend = BlobBackend(self.binding, self.wake, BlobConfig(key_prefix=KEY_PREFIX)) + self.threads = set(threading.enumerate()) - before + + def rows(self, names: Sequence[str], layer_group_index: int = 0) -> Rows: + keys = np.stack([object_key(name) for name in names]) if names else _no_keys() + slots = self.region.slots(*(self.names.index(name) for name in names)) + return Rows(layer_group_index, keys, slots, 0) + + def view(self, *names: str) -> CpuView: + return CpuView([self.rows(names)]) + + def write(self, name: str, data: bytes) -> None: + self.region.write(self.names.index(name), data) + + def read(self, name: str) -> bytes: + return self.region.read(self.names.index(name)) + + def wait_outcome(self, attempt: BlobAttempt) -> Outcome: + self.wake.wait_for(lambda: attempt.poll() is not None, f"an outcome for {attempt.label}") + return attempt.poll() + + def op(self, name: str) -> OpHandle: + """The latest driver operation carrying the row ``name``, once ``put`` or ``get`` returned + it.""" + return self.driver.op_with(store_key(name)) + + def wait_quiesced(self, attempt: BlobAttempt) -> None: + self.wake.wait_for(attempt.quiesced, f"{attempt.label} quiescing") + + def close(self) -> None: + assert self.backend.stats()["attempts_active"] == 0, "an attempt outlived its transfers" + labels = self.backend.shutdown(WAIT_S) + assert labels == frozenset(), f"still accessible after shutdown: {sorted(labels)}" + self.binding.release_hold() + assert self.hold.release_threads == [threading.get_ident()] + for thread in self.threads: + thread.join(WAIT_S) + assert not any(thread.is_alive() for thread in self.threads) + # A background ``MemoryDriver``'s threads count as finished a moment before they exit. + wait_until( + lambda: not any(t.name == "memory-driver-op" for t in threading.enumerate()), + "the driver's operation threads to exit", + ) + # The owner that opened the driver closes it, last: after the drain, the revocations and + # the hold's return. The backend never does. + self.driver.close() + ending = [call.method for call in self.driver.calls if call.method in ("drain", "close")] + assert ending.count("close") == 1 and ending[-1] == "close" and "drain" in ending, ending + + +def _no_keys() -> np.ndarray: + return np.zeros((0, OBJECT_KEY_WIDTH), dtype=np.uint8) + + +class FakeLeaseHolder: + """The runtime's side of an attempt, reduced to what a backend can see of it. + + One lease on the endpoint's slots per submission; one decision per attempt, taken from the + outcome once it landed or from nothing at the holder's own deadline, and never revisited; the + lease given back only once the attempt quiesced, whatever was decided. The region hold is the + manager's: it goes back through the binding, on the manager's thread, once nothing is + registered any more. ``on_wake`` is the holder's reaction to a wake; a test calls the steps + one by one when it wants them apart. + """ + + def __init__(self, endpoint: Endpoint) -> None: + self.endpoint = endpoint + self.attempts: dict[str, BlobAttempt] = {} + self.decisions: dict[str, Outcome | None] = {} # None: the deadline came first + self.leased: set[str] = set() # attempts whose slots the lease still covers + + def fetch(self, *names: str, hint: Hint = None) -> BlobAttempt: + return self._submit(self.endpoint.backend.fetch(self.endpoint.view(*names), hint)) + + def publish(self, *names: str) -> BlobAttempt: + return self._submit(self.endpoint.backend.publish(self.endpoint.view(*names), None)) + + def decide(self, attempt: BlobAttempt) -> None: + """Takes the outcome as the decision, or nothing when none has landed, as at the + holder's deadline. The first decision stands.""" + self.decisions.setdefault(attempt.label, attempt.poll()) + + def delivered(self, attempt: BlobAttempt) -> list[list[bool]] | None: + """What the decision counts as delivered; ``None`` when it was taken from nothing.""" + outcome = self.decisions[attempt.label] + return None if outcome is None else masks(outcome) + + def return_lease_if_quiesced(self, attempt: BlobAttempt) -> bool: + if attempt.quiesced(): + self.leased.discard(attempt.label) + return attempt.label not in self.leased + + def on_wake(self) -> None: + for attempt in self.attempts.values(): + if attempt.poll() is not None: + self.decide(attempt) + self.return_lease_if_quiesced(attempt) + + def shut_down(self, budget: float) -> frozenset[str]: + """The manager's shutdown: what the backend still holds, and the hold given back only + when that is nothing. Called again later when it was not.""" + labels = self.endpoint.backend.shutdown(budget) + if not self.endpoint.binding.live_labels(): + self.endpoint.binding.release_hold() + return labels + + def _submit(self, attempt: BlobAttempt) -> BlobAttempt: + self.attempts[attempt.label] = attempt + self.leased.add(attempt.label) + return attempt diff --git a/tests/unittest/_torch/disaggregation/blob/test_blob_backend.py b/tests/unittest/_torch/disaggregation/blob/test_blob_backend.py new file mode 100644 index 000000000000..5359a43c0f78 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_blob_backend.py @@ -0,0 +1,1768 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The blob backend over the in-memory driver. + +Covers the store's behaviour cases (a late write after a fault, results and access end in +either order, cancel and shutdown races) and the backend's own rules: one driver call per layer +group, a submission returns while its operations run, only regions the driver can access are +registered and rows elsewhere are refused, each operation's access end is counted once, a fault +waits for every call, every call reaches the driver as it is submitted and a stopped attempt's +calls are cancelled there once each, a call is queried only when its driver calls its on_change, +a driver that breaks its protocol gets no more work, and shutdown never calls the driver. Every +wait is bounded; a sleep of PROMPT_S only stands for "nothing happened meanwhile". +""" + +from __future__ import annotations + +import math +import threading +import time +from collections.abc import Callable +from dataclasses import replace + +import numpy as np +import pytest +from cpu_fixtures import ( + PROMPT_S, + SLOT_BYTES, + WAIT_S, + CpuRegion, + CpuView, + Endpoint, + FakeHold, + KeyGatedStore, + RollbackFailsDriver, + Rows, + TogetherDriverSpy, + commit_object, + masks, + object_key, + store_key, + wait_until, +) +from together_driver import TogetherDriver + +from tensorrt_llm._torch.disaggregation.base.backend import ( + Capabilities, + InvalidSubmission, + MemoryKind, + SubmissionRejected, +) +from tensorrt_llm._torch.disaggregation.blob import backend as backend_module +from tensorrt_llm._torch.disaggregation.blob.backend import ( + SHUTDOWN_FAULT, + BlobBackend, + BlobConfig, + ConstructionRejected, + RegionBinding, +) +from tensorrt_llm._torch.disaggregation.blob.store import ( + DriverRejected, + RowBatch, + RowIO, + Span, + StoreError, +) +from tensorrt_llm._torch.disaggregation.testing.memory_store import ( + DEFAULT_CAPABILITIES, + Commit, + MemoryDriver, + MemoryOp, + MemoryStore, +) + +# An exception that escapes a thread fails the test. +pytestmark = [ + pytest.mark.cpu_only, + pytest.mark.filterwarnings("error::pytest.PytestUnhandledThreadExceptionWarning"), +] + +# A test that leaves a call whose access never ends keeps that backend's progress thread alive, as +# it would in production; the thread-leak check would report it. +keeps_backend_threads = pytest.mark.threadleak(enabled=False) + + +@pytest.fixture +def store() -> MemoryStore: + return MemoryStore() + + +def commit_objects(store: MemoryStore, rows: dict[str, bytes]) -> None: + """Commits ``rows`` straight into ``store`` under the keys the backend derives.""" + for name, data in rows.items(): + assert commit_object(store, name, data) is Commit.STORED + + +def store_calls(endpoint: Endpoint) -> list[str]: + return [call.method for call in endpoint.driver.calls if call.method in ("get", "put")] + + +ONE_ROW_PER_CALL = replace(DEFAULT_CAPABILITIES, max_rows_per_op=1) + + +def one_row_per_call(store: MemoryStore, **options) -> MemoryDriver: + """A driver that takes one row per operation, so each row of a group is a call of its own.""" + return MemoryDriver(store, capabilities=ONE_ROW_PER_CALL, **options) + + +# Results. + + +@pytest.mark.parametrize("case", ["fetch_not_ready", "publish_declined"]) +def test_not_ready_and_declined_rows_are_undelivered_without_a_fault(case: str) -> None: + # A fetch finds the first object still being written (NOT_READY); a publish finds no room + # after the first row (DECLINED). + store = MemoryStore(capacity=1 if case == "publish_declined" else None) + endpoint = Endpoint(one_row_per_call(store), ["A", "B", "C"]) + if case == "fetch_not_ready": + commit_objects(store, {"B": b"bbbbbbbb", "C": b"cccccccc"}) + store.begin_write(store_key("A")) + attempt = endpoint.backend.fetch(endpoint.view("A", "B", "C"), None) + expected = [False, True, True] + else: + for name in "ABC": + endpoint.write(name, name.encode() * SLOT_BYTES) + attempt = endpoint.backend.publish(endpoint.view("A", "B", "C"), None) + expected = [True, False, False] + outcome = endpoint.wait_outcome(attempt) + assert outcome.fault is None and masks(outcome) == [expected] + assert len(store_calls(endpoint)) == 3 # a skipped row stops nothing + endpoint.wait_quiesced(attempt) + endpoint.close() + + +def test_masks_follow_view_group_order(store: MemoryStore) -> None: + commit_objects(store, {"A": b"aaaaaaaa", "C": b"cccccccc", "E": b"eeeeeeee"}) + reader = Endpoint(MemoryDriver(store), ["A", "B", "C", "D", "E"]) + view = CpuView([reader.rows(["C", "D"], 2), reader.rows(["A"], 0), reader.rows(["E", "B"], 1)]) + attempt = reader.backend.fetch(view, None) + outcome = reader.wait_outcome(attempt) + assert masks(outcome) == [[True, False], [True], [True, False]] + outcome.check_mask_shapes(view) + assert [len(call.rows) for call in reader.driver.calls_to("get")] == [2, 1, 2] + reader.wait_quiesced(attempt) + reader.close() + + +def test_empty_view_concludes_and_quiesces_at_once(store: MemoryStore) -> None: + reader = Endpoint(MemoryDriver(store), ["A"]) + view = CpuView([reader.rows([], 0), reader.rows([], 1)]) + attempt = reader.backend.fetch(view, None) + outcome = attempt.poll() + assert outcome is not None and outcome.fault is None + assert [mask.shape for mask in outcome.delivered_masks] == [(0,), (0,)] + assert attempt.quiesced() + assert reader.wake.count >= 1 + assert store_calls(reader) == [] + reader.close() + + +@pytest.mark.parametrize("max_rows_per_op", [None, 4], ids=["no_limit", "driver_limit"]) +def test_each_group_goes_in_one_call_split_only_by_the_driver_limit( + store: MemoryStore, max_rows_per_op: int | None +) -> None: + names = [f"R{i}" for i in range(12)] + capabilities = replace(DEFAULT_CAPABILITIES, max_rows_per_op=max_rows_per_op) + writer = Endpoint(MemoryDriver(store, capabilities=capabilities), names) + view = CpuView([writer.rows(names[:10], 0), writer.rows(names[10:], 1)]) + attempt = writer.backend.publish(view, None) + assert masks(writer.wait_outcome(attempt)) == [[True] * 10, [True] * 2] + expected = [10, 2] if max_rows_per_op is None else [4, 4, 2, 2] + assert [len(call.rows) for call in writer.driver.calls_to("put")] == expected + limit = None if max_rows_per_op is None else (4, 4 * SLOT_BYTES) + assert writer.backend.driver_call_limit(0) == limit + writer.wait_quiesced(attempt) + writer.close() + + +def test_capabilities_are_the_driver_directions(store: MemoryStore) -> None: + driver = MemoryDriver(store, capabilities=replace(DEFAULT_CAPABILITIES, get=False)) + writer = Endpoint(driver, ["A"]) + assert writer.backend.capabilities == Capabilities(fetch=False, publish=True, cancel=True) + with pytest.raises(InvalidSubmission, match="cannot get"): + writer.backend.fetch(writer.view("A"), None) + assert store_calls(writer) == [] + attempt = writer.backend.publish(writer.view("A"), None) + assert masks(writer.wait_outcome(attempt)) == [[True]] + writer.wait_quiesced(attempt) + writer.close() + + +def test_blob_config_holds_only_the_key_prefix() -> None: + assert BlobConfig() == BlobConfig(key_prefix="trtllm-kv") + for prefix in ("", "a/b", "a b", None): + with pytest.raises(ValueError, match="prefix"): + BlobConfig(key_prefix=prefix) + # Tuning such as workers belongs to the driver's own section. + with pytest.raises(TypeError, match="workers"): + BlobConfig(workers=2) + + +# Results, access end and late writes. + + +def test_late_write_after_fault_changes_nothing(store: MemoryStore) -> None: + commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = one_row_per_call(store, controlled=True) + reader = Endpoint(driver, ["A", "B"], fill=0x2D) + attempt = reader.backend.fetch(reader.view("A", "B"), None) + op_a, op_b = reader.op("A"), reader.op("B") + driver.complete(op_a) + driver.end_access(op_a) + driver.fail(op_b, "timeout") + outcome = reader.wait_outcome(attempt) + assert masks(outcome) == [[True, False]] + assert "timeout" in outcome.fault + assert not attempt.quiesced() + + driver.write_late(op_b, 0, b"bbbbbbbb") + assert reader.read("B") == b"bbbbbbbb" + assert attempt.poll() is outcome + driver.end_access(op_b) + reader.wait_quiesced(attempt) + assert attempt.poll() is outcome + reader.close() + + +def test_access_end_before_result_without_touching_the_view(store: MemoryStore) -> None: + commit_objects(store, {"A": b"aaaaaaaa"}) + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A"], fill=0x2D) + view = reader.view("A") + attempt = reader.backend.fetch(view, None) + view.seal() + op = reader.op("A") + driver.transfer(op) + driver.end_access(op) + reader.wait_quiesced(attempt) + assert attempt.poll() is None + driver.publish_results(op) + outcome = reader.wait_outcome(attempt) + assert masks(outcome) == [[True]] + assert reader.read("A") == b"aaaaaaaa" + reader.close() + + +@pytest.mark.parametrize("first_row", ["ready", "faulty"]) +def test_a_later_call_fault_waits_for_earlier_calls_and_the_first_recorded_is_reported( + first_row: str, +) -> None: + store = KeyGatedStore(gets=True, keys=[store_key("A")]) + a = b"aaaaaaaa" if first_row == "ready" else b"aaa" # the wrong size: a second fault + commit_objects(store, {"A": a, "B": b"short"}) + driver = one_row_per_call(store, background=True) + # B's fault stops the fetch and forwards a cancel; ignoring it lets A still reach the gate. + driver.set_cancel_effective(False) + reader = Endpoint(driver, ["A", "B"]) + attempt = reader.backend.fetch(reader.view("A", "B"), None) + try: + store.wait_for_calls_at_gate(1) + # B's fault is recorded while A is still in the store. + wait_until(lambda: reader.backend.stats()["calls_awaiting_results"] == 1, "B's fault") + assert attempt.poll() is None + assert attempt.quiesced() is False + finally: + store.gate.set() + outcome = reader.wait_outcome(attempt) + assert masks(outcome) == [[first_row == "ready", False]] # A's row is kept for reuse + assert "row 1: object is 5 bytes" in outcome.fault and "3 bytes" not in outcome.fault + reader.wait_quiesced(attempt) + reader.close() + + +def three_rows(store: MemoryStore, objects_in_store: bool, **options) -> Endpoint: + """An endpoint whose rows A, B and C each go in a call of their own, handed over in that order, + with their bytes in the store when ``objects_in_store`` and in the slots otherwise. ``options`` + go to the ``MemoryDriver``.""" + data = {name: name.encode() * SLOT_BYTES for name in "ABC"} + if objects_in_store: + commit_objects(store, data) + endpoint = Endpoint(one_row_per_call(store, **options), list(data)) + if not objects_in_store: + for name, row in data.items(): + endpoint.write(name, row) + return endpoint + + +def test_a_row_fault_stops_a_fetch(store: MemoryStore) -> None: + reader = three_rows(store, objects_in_store=True, controlled=True) + driver = reader.driver.inner + attempt = reader.backend.fetch(reader.view("A", "B", "C"), None) + op_a, op_b, op_c = (reader.op(name) for name in "ABC") + driver.fail(op_a, "row lost") + # A prefix reuse has no use for the rows past the first gap, so B's and C's calls, which the + # driver still holds, are cancelled. + reader.driver.wait_for(lambda: len(reader.driver.calls_to("cancel")) == 2, "two cancels") + cancelled = [c.op for c in reader.driver.calls_to("cancel")] + assert cancelled[0] is op_b and cancelled[1] is op_c + outcome = reader.wait_outcome(attempt) + assert outcome.fault == "group 0 row 0: row lost" + assert masks(outcome) == [[False, False, False]] + assert store_calls(reader) == ["get"] * 3 # every call reached the driver as submitted + driver.end_access(op_a) # B and C ended theirs as they were cancelled + reader.wait_quiesced(attempt) + reader.close() + + +def test_a_row_fault_leaves_the_rest_of_a_publish_to_go_on(store: MemoryStore) -> None: + writer = three_rows(store, objects_in_store=False, controlled=True) + driver = writer.driver.inner + attempt = writer.backend.publish(writer.view("A", "B", "C"), None) + op_a, op_b, op_c = (writer.op(name) for name in "ABC") + driver.fail(op_a, "row lost") + time.sleep(PROMPT_S) # a cancel would have been forwarded by now + assert writer.driver.calls_to("cancel") == [] + for op in (op_b, op_c): + driver.complete(op) + outcome = writer.wait_outcome(attempt) + assert outcome.fault == "group 0 row 0: row lost" + assert masks(outcome) == [[False, True, True]] # a publish's rows stand alone + for op in (op_a, op_b, op_c): + driver.end_access(op) + writer.wait_quiesced(attempt) + writer.close() + + +class RefusesItsSecondPut(MemoryDriver): + puts = 0 + + def put(self, rows, *, on_change=None): + self.puts += 1 + if self.puts == 2: + raise DriverRejected("full") + return super().put(rows, on_change=on_change) + + +def test_a_call_that_concludes_without_a_fault_stops_none_of_the_others(store: MemoryStore) -> None: + commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = one_row_per_call(store, controlled=True) + reader = Endpoint(driver, ["A", "B"]) + attempt = reader.backend.fetch(reader.view("A", "B"), None) + op_a, op_b = reader.op("A"), reader.op("B") + driver.complete(op_a) + driver.end_access(op_a) + wait_until(lambda: reader.backend.stats()["calls_awaiting_results"] == 1, "A's results") + time.sleep(PROMPT_S) # a cancel would have been forwarded by now + assert reader.driver.calls_to("cancel") == [] + driver.complete(op_b) + driver.end_access(op_b) + assert masks(reader.wait_outcome(attempt)) == [[True, True]] + reader.wait_quiesced(attempt) + reader.close() + + +def test_a_refused_put_cancels_the_calls_of_its_publish_already_in_the_driver( + store: MemoryStore, +) -> None: + driver = RefusesItsSecondPut(store, controlled=True, capabilities=ONE_ROW_PER_CALL) + writer = Endpoint(driver, ["A", "B", "C"]) + for name in "ABC": + writer.write(name, name.encode() * SLOT_BYTES) + attempt = writer.backend.publish(writer.view("A", "B", "C"), None) + op_a = writer.op("A") + writer.driver.wait_for(lambda: writer.driver.calls_to("cancel", op_a), "A's cancel") + outcome = writer.wait_outcome(attempt) # A cancelled, B refused, C never handed over + assert masks(outcome) == [[False, False, False]] + assert outcome.fault == "group 0 row 1: refused: full" + writer.wait_quiesced(attempt) + writer.close() + + +def test_a_refused_put_stops_its_publish_but_not_the_backend(store: MemoryStore) -> None: + writer = three_rows(store, objects_in_store=False) + writer.driver.fail("put", DriverRejected("busy")) + attempt = writer.backend.publish(writer.view("A", "B", "C"), None) + outcome = writer.wait_outcome(attempt) + assert outcome.fault == "group 0 row 0: refused: busy" + assert masks(outcome) == [[False, False, False]] # the call failed as a whole + assert store_calls(writer) == ["put"] + writer.wait_quiesced(attempt) + # A refusal breaks no protocol, so the backend takes the next submission. + again = writer.backend.publish(writer.view("A", "B", "C"), None) + assert masks(writer.wait_outcome(again)) == [[True, True, True]] + writer.wait_quiesced(again) + writer.close() + + +def test_an_outcome_waits_for_every_row_of_a_call_that_concludes_in_parts( + store: MemoryStore, +) -> None: + commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A", "B"]) + attempt = reader.backend.fetch(reader.view("A", "B"), None) + op = reader.op("A") # one call holds both rows + driver.complete(op, [0]) # tells of A; B is still pending + assert attempt.poll() is None + driver.complete(op, [1]) + driver.end_access(op) + assert masks(reader.wait_outcome(attempt)) == [[True, True]] + reader.wait_quiesced(attempt) + reader.close() + + +@pytest.mark.parametrize("first_half", ["results_first", "access_end_first"]) +def test_results_and_access_end_are_each_recorded_once_per_call( + store: MemoryStore, first_half: str +) -> None: + commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = one_row_per_call(store, controlled=True) + reader = Endpoint(driver, ["A", "B"]) + attempt = reader.backend.fetch(reader.view("A", "B"), None) + op_a, op_b = reader.op("A"), reader.op("B") + # A lands one half only, and the driver tells of it again and again. + if first_half == "results_first": + driver.complete(op_a) + else: + driver.transfer(op_a) + driver.end_access(op_a) + for _ in range(5): + reader.driver.call_on_change(op_a) + # Recording A's half on every on_change would already have concluded the attempt (results_first) + # or quiesced it (access_end_first). + assert attempt.poll() is None # B has no result yet + assert not attempt.quiesced() # B's access is open + driver.complete(op_b) + driver.end_access(op_b) + if first_half == "results_first": + driver.end_access(op_a) + else: + driver.publish_results(op_a) + assert masks(reader.wait_outcome(attempt)) == [[True, True]] + reader.wait_quiesced(attempt) + reader.close() + + +@keeps_backend_threads +@pytest.mark.parametrize("ending", ["access_ended", "unprovable"]) +def test_news_of_a_settled_operation_changes_nothing(store: MemoryStore, ending: str) -> None: + commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A", "B"], fill=0x2D) + first = reader.backend.fetch(reader.view("A"), None) + op_a = reader.op("A") + if ending == "access_ended": + driver.complete(op_a) + driver.end_access(op_a) + reader.wait_outcome(first) # quiescing alone may come before the results are recorded + reader.wait_quiesced(first) + else: + driver.fail(op_a, "the store gave up") + driver.mark_unprovable(op_a) # nothing the driver could wait on was issued + assert "the store gave up" in reader.wait_outcome(first).fault + wait_until(lambda: reader.backend.stats()["calls_unprovable"] == 1, "the count") + outcome, wakes = first.poll(), reader.wake.count + reader.driver.call_on_change(op_a) # a driver may call on_change more than once + assert first.poll() is outcome and reader.wake.count == wakes + + second = reader.backend.fetch(reader.view("B"), None) # the backend serves on + op_b = reader.op("B") + driver.complete(op_b) + driver.end_access(op_b) + reader.wait_outcome(second) + reader.wait_quiesced(second) + if ending == "access_ended": + reader.close() + return + assert not first.quiesced() # its slots stay held + assert reader.backend.stats()["calls_awaiting_access_end"] == 1 + assert reader.backend.shutdown(0) == frozenset({first.label, "region:0"}) + + +@pytest.mark.parametrize( + "method", ["drain", "unregister"], ids=["drain_answers_false", "unregister_raises"] +) +def test_the_progress_thread_retries_drain_and_revocation_and_the_hold_stays( + store: MemoryStore, method: str +) -> None: + endpoint = Endpoint(MemoryDriver(store), ["A"]) + if method == "drain": + endpoint.driver.answer("drain", False, times=1_000_000) + else: + endpoint.driver.fail("unregister", StoreError("busy"), times=1_000_000) + assert endpoint.backend.shutdown(0) == frozenset({"region:0"}) + endpoint.driver.wait_for(lambda: len(endpoint.driver.calls_to(method)) >= 3, "retries") + if method == "drain": + assert endpoint.driver.calls_to("unregister") == [] # never revoked before a True drain + with pytest.raises(RuntimeError, match="region:0"): + endpoint.binding.release_hold() + assert endpoint.binding.live_labels() == frozenset({"region:0"}) + if method == "drain": + endpoint.driver.answer("drain", False, times=0) + else: + endpoint.driver.fail("unregister", StoreError("busy"), times=0) + # Nothing else calls in: the progress thread retries on its own. + wait_until(lambda: not endpoint.binding.live_labels(), "the revocation to succeed") + endpoint.close() + + +# Cancel. + + +@pytest.mark.parametrize("first", ["cancel_first", "complete_first"]) +def test_cancel_racing_completion_gives_one_outcome(store: MemoryStore, first: str) -> None: + commit_objects(store, {"A": b"aaaaaaaa"}) + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A"]) + attempt = reader.backend.fetch(reader.view("A"), None) + op = reader.op("A") + if first == "cancel_first": + attempt.cancel() + assert attempt.poll() is None + assert reader.wake.count == 0 # nothing landed, so nothing to wake for + # Forwarded: the driver cancels the row it had not started, which touched nothing. + reader.driver.wait_for(lambda: reader.driver.calls_to("cancel", op), "the cancel") + expected = [[False]] + else: + driver.complete(op) + driver.end_access(op) + reader.wait_outcome(attempt) + reader.wait_quiesced(attempt) + wakes = reader.wake.count + attempt.cancel() + assert reader.wake.count == wakes + expected = [[True]] + outcome = reader.wait_outcome(attempt) + reader.wait_quiesced(attempt) + assert attempt.poll() is outcome + assert masks(outcome) == expected and outcome.fault is None + assert 1 <= reader.wake.count <= 2 # the outcome and the quiescing, together or apart + # Only a call the driver still holds is cancelled, once. + assert len(reader.driver.calls_to("cancel")) == (first == "cancel_first") + reader.close() + + +def test_a_cancel_reaches_the_driver_promptly_and_once(store: MemoryStore) -> None: + commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = MemoryDriver(store, controlled=True) + driver.set_cancel_effective(False) # the calls stay open + reader = Endpoint(driver, ["A", "B"]) + attempt = reader.backend.fetch(reader.view("A"), None) + op = reader.op("A") + time.sleep(0.05) # lets the progress thread settle into its wait + attempt.cancel() + started = time.monotonic() + reader.driver.wait_for(lambda: reader.driver.calls_to("cancel", op), "the cancel") + assert time.monotonic() - started < PROMPT_S # woken by the cancel's poke + attempt.cancel() # again: nothing more to forward + # Another attempt's cancel runs another progress round while A's call is still open. + other = reader.backend.fetch(reader.view("B"), None) + other_op = reader.op("B") + other.cancel() + reader.driver.wait_for(lambda: reader.driver.calls_to("cancel", other_op), "B's cancel") + assert len(reader.driver.calls_to("cancel", op)) == 1 + for each in (op, other_op): + driver.complete(each) + driver.end_access(each) + reader.wait_quiesced(attempt) + reader.wait_quiesced(other) + reader.close() + + +class FaultsItsFirstGetInside(MemoryDriver): + """Faults its first get inside ``get``, telling of it there, and gives the progress thread time + for a round before ``get`` returns: the stop lands before any handle is known.""" + + gets = 0 + + def get(self, rows, hint=None, *, on_change=None): + op = super().get(rows, hint, on_change=on_change) + self.gets += 1 + if self.gets == 1: + self.fail(op, "row lost") + self.end_access(op) + time.sleep(0.05) + return op + + +def test_a_stop_during_the_hand_over_reaches_the_calls_handed_over_after_it( + store: MemoryStore, +) -> None: + commit_objects(store, {name: name.encode() * SLOT_BYTES for name in "ABC"}) + driver = FaultsItsFirstGetInside(store, controlled=True, capabilities=ONE_ROW_PER_CALL) + reader = Endpoint(driver, ["A", "B", "C"]) + attempt = reader.backend.fetch(reader.view("A", "B", "C"), None) + outcome = reader.wait_outcome(attempt) # B's and C's calls conclude only once cancelled + assert masks(outcome) == [[False, False, False]] + assert outcome.fault == "group 0 row 0: row lost" + reader.wait_quiesced(attempt) + reader.close() + + +def test_a_cancel_asked_while_the_progress_thread_is_cancelling_is_still_forwarded( + store: MemoryStore, +) -> None: + commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A", "B"]) + first = reader.backend.fetch(reader.view("A"), None) + second = reader.backend.fetch(reader.view("B"), None) + op_b = reader.op("B") + gate = reader.driver.gate("cancel") + first.cancel() + reader.driver.wait_for(lambda: reader.driver.calls_to("cancel"), "A's cancel") + second.cancel() # the progress thread is inside the driver's cancel of A's call + gate.set() + reader.driver.wait_for(lambda: reader.driver.calls_to("cancel", op_b), "B's cancel") + for attempt in (first, second): + reader.wait_outcome(attempt) + reader.wait_quiesced(attempt) + reader.close() + + +def test_a_driver_without_cancel_is_never_asked_to(store: MemoryStore) -> None: + commit_objects(store, {"B": b"bbbbbbbb"}) # A is missing: the fetch stops at its fault + capabilities = replace(DEFAULT_CAPABILITIES, cancel=False, max_rows_per_op=1) + driver = MemoryDriver(store, controlled=True, capabilities=capabilities) + reader = Endpoint(driver, ["A", "B"]) + assert not reader.backend.capabilities.cancel # every call is in the driver at once + attempt = reader.backend.fetch(reader.view("A", "B"), None) + op_a, op_b = reader.op("A"), reader.op("B") + driver.fail(op_a, "lost") + driver.end_access(op_a) + time.sleep(PROMPT_S) # a cancel would have been forwarded by now + driver.complete(op_b) # it ran on: nobody could cancel it + driver.end_access(op_b) + outcome = reader.wait_outcome(attempt) + assert masks(outcome) == [[False, True]] and "lost" in outcome.fault + reader.wait_quiesced(attempt) + assert reader.driver.calls_to("cancel") == [] + reader.close() + + +# The driver tells of changes through on_change. + + +def test_a_call_is_queried_only_when_its_driver_calls_on_change(store: MemoryStore) -> None: + commit_objects(store, {"A": b"aaaaaaaa"}) + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A"]) + attempt = reader.backend.fetch(reader.view("A"), None) + op = reader.op("A") + time.sleep(PROMPT_S) # a polling backend would have queried several times by now + assert reader.driver.calls_to("poll", op) == [] + driver.complete(op) # tells + assert masks(reader.wait_outcome(attempt)) == [[True]] + driver.end_access(op) # tells + reader.wait_quiesced(attempt) + # Queried on the thread that called on_change: here the test's own. + queries = reader.driver.calls_to("poll", op) + assert len(queries) == 2 and all(c.thread == threading.current_thread().name for c in queries) + reader.close() + + +def test_a_teardown_retries_the_drain_with_backoff(store: MemoryStore) -> None: + commit_objects(store, {"B": b"bbbbbbbb"}) + driver = MemoryDriver(store, controlled=True) + endpoint = Endpoint(driver, ["A", "B"]) + # An operation outside the backend, already started: no drain can end it. + row = RowIO(store_key("B"), (Span(endpoint.region.slot(1), SLOT_BYTES),)) + outside = driver.get(RowBatch.of([row])) + driver.transfer(outside) + assert endpoint.backend.shutdown(0) == frozenset({"region:0"}) + time.sleep(PROMPT_S) # polled with backoff: about ten drains + assert len(endpoint.driver.calls_to("drain")) < 30 + driver.publish_results(outside) + driver.end_access(outside) + wait_until(lambda: not endpoint.binding.live_labels(), "the revocation") + endpoint.close() + + +def test_a_teardown_follows_the_last_attempt_to_quiesce( + store: MemoryStore, monkeypatch: pytest.MonkeyPatch +) -> None: + # Retries far apart: only the quiescing itself can bring the teardown on in time. + monkeypatch.setattr(backend_module, "TEARDOWN_RETRY_MIN_S", 10 * WAIT_S) + monkeypatch.setattr(backend_module, "TEARDOWN_RETRY_MAX_S", 10 * WAIT_S) + commit_objects(store, {"A": b"aaaaaaaa"}) + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A"]) + attempt = reader.backend.fetch(reader.view("A"), None) + op = reader.op("A") + driver.transfer(op) # running: no drain can end it + assert reader.backend.shutdown(0) == frozenset({attempt.label, "region:0"}) + reader.driver.wait_for(lambda: reader.driver.calls_to("drain"), "the first teardown round") + driver.publish_results(op) + driver.end_access(op) + wait_until(lambda: not reader.binding.live_labels(), "the revocation") + reader.close() + + +def test_a_change_told_before_put_or_get_returns_is_not_lost(store: MemoryStore) -> None: + commit_objects(store, {"A": b"aaaaaaaa"}) + # An automatic driver runs each operation, and tells of it, inside the put or get that takes + # it: before the backend has the operation's handle. + endpoint = Endpoint(MemoryDriver(store), ["A", "B"]) + endpoint.write("B", b"bbbbbbbb") + fetched = endpoint.backend.fetch(endpoint.view("A"), None) + offered = endpoint.backend.publish(endpoint.view("B"), None) + for attempt in (fetched, offered): + assert masks(endpoint.wait_outcome(attempt)) == [[True]] + endpoint.wait_quiesced(attempt) + endpoint.close() + + +# Driver failures. + + +def test_a_refusal_faults_the_call_and_ends_its_access(store: MemoryStore) -> None: + writer = Endpoint(MemoryDriver(store), ["A"]) + writer.driver.fail("put", DriverRejected("x")) + attempt = writer.backend.publish(writer.view("A"), None) + outcome = writer.wait_outcome(attempt) + assert masks(outcome) == [[False]] and outcome.fault == "group 0 row 0: refused: x" + writer.wait_quiesced(attempt) # nothing escaped + assert writer.backend.stats()["calls_unprovable"] == 0 + writer.close() + + +@keeps_backend_threads +def test_any_other_raise_from_the_driver_holds_the_call_and_refuses_later_work( + store: MemoryStore, +) -> None: + writer = Endpoint(MemoryDriver(store), ["A"]) + writer.driver.fail("put", RuntimeError("x")) + attempt = writer.backend.publish(writer.view("A"), None) + outcome = writer.wait_outcome(attempt) + assert masks(outcome) == [[False]] + assert outcome.fault == "group 0 row 0: the store driver raised RuntimeError('x')" + assert not attempt.quiesced() # the rows may be accessed for good + assert writer.backend.stats()["calls_unprovable"] == 1 + with pytest.raises(SubmissionRejected, match="broke its protocol: the store driver raised"): + writer.backend.publish(writer.view("A"), None) + assert attempt.label in writer.backend.shutdown(0) + with pytest.raises(RuntimeError, match="region:0"): + writer.binding.release_hold() + assert not writer.hold.released + + +class NoHandleDriver(MemoryDriver): + """Answers ``put`` with None in place of a handle.""" + + def put(self, rows, *, on_change=None): + return None + + +@keeps_backend_threads +def test_a_put_that_returns_no_handle_holds_the_call_and_refuses_later_work( + store: MemoryStore, +) -> None: + writer = Endpoint(NoHandleDriver(store), ["A"]) + attempt = writer.backend.publish(writer.view("A"), None) + outcome = writer.wait_outcome(attempt) + assert masks(outcome) == [[False]] + assert outcome.fault == "group 0 row 0: the store driver returned no handle" + assert not attempt.quiesced() # nothing proves the driver stopped reading the slot + assert writer.backend.stats()["calls_unprovable"] == 1 + with pytest.raises(SubmissionRejected, match="broke its protocol: .* returned no handle"): + writer.backend.publish(writer.view("A"), None) + assert writer.backend.shutdown(0) == frozenset({attempt.label, "region:0"}) + + +@keeps_backend_threads +def test_a_protocol_violation_leaves_calls_already_in_the_driver_alone() -> None: + store = KeyGatedStore(gets=True) + commit_objects(store, {"X": b"xxxxxxxx"}) + endpoint = Endpoint(MemoryDriver(store, background=True), ["X", "Y"]) + endpoint.write("Y", b"yyyyyyyy") + in_driver = endpoint.backend.fetch(endpoint.view("X"), None) + try: + store.wait_for_calls_at_gate(1) # its only call is inside the store + endpoint.driver.fail("put", RuntimeError("lost")) + broken = endpoint.backend.publish(endpoint.view("Y"), None) + assert "the store driver raised" in endpoint.wait_outcome(broken).fault + finally: + store.gate.set() + outcome = endpoint.wait_outcome(in_driver) + assert outcome.fault is None and masks(outcome) == [[True]] + endpoint.wait_quiesced(in_driver) + assert endpoint.backend.shutdown(0) == frozenset({broken.label, "region:0"}) + + +@keeps_backend_threads +def test_a_driver_raise_stops_the_rest_of_its_submission(store: MemoryStore) -> None: + writer = three_rows(store, objects_in_store=False) + writer.driver.fail("put", RuntimeError("lost")) + broken = writer.backend.publish(writer.view("A", "B", "C"), None) + outcome = writer.wait_outcome(broken) + assert outcome.fault == "group 0 row 0: the store driver raised " + repr(RuntimeError("lost")) + assert masks(outcome) == [[False, False, False]] + assert store_calls(writer) == ["put"] # the raise stopped the publish: B and C never went + assert writer.backend.stats()["calls_awaiting_access_end"] == 1 # only A's + assert writer.backend.shutdown(0) == frozenset({broken.label, "region:0"}) + + +class ShortPollDriver(MemoryDriver): + """Once ``short`` is set, answers the next poll of a multi-row operation one result short.""" + + short = False + + def poll(self, op): + results = super().poll(op) + if self.short and len(results) > 1: + self.short = False + return results[:-1] + return results + + +@keeps_backend_threads +def test_a_poll_with_the_wrong_row_count_faults_the_call_and_keeps_it_held( + store: MemoryStore, +) -> None: + commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = ShortPollDriver(store) + driver.short = True # the first query of the operation, as get calls on_change + reader = Endpoint(driver, ["A", "B", "C"]) + reader.write("C", b"cccccccc") + attempt = reader.backend.fetch(reader.view("A", "B"), None) + outcome = reader.wait_outcome(attempt) + assert masks(outcome) == [[False, False]] + assert "gave 1 results for 2 rows" in outcome.fault + # Later answers have the right length and say access ended; after the violation the call is + # unprovable instead: counted, its slots held. + assert reader.backend.stats()["calls_unprovable"] == 1 + reader.driver.call_on_change(reader.op("A")) + assert not attempt.quiesced() + # The backend survived the bad answer and refuses all later work, naming why. + assert all(thread.is_alive() for thread in reader.threads) + with pytest.raises(SubmissionRejected, match="broke its protocol: .* gave 1 results for 2"): + reader.backend.publish(reader.view("C"), None) + assert store_calls(reader) == ["get"] + assert reader.backend.shutdown(0) == frozenset({attempt.label, "region:0"}) + # The driver's own operation ended, so its drain answers True; the call the backend holds + # still keeps the region registered. + reader.driver.wait_for(lambda: reader.driver.calls_to("drain"), "a drain") + time.sleep(PROMPT_S) # a revocation would have come by now + assert driver.live_registrations() + + +@keeps_backend_threads +def test_the_calls_after_a_protocol_violation_never_reach_the_driver( + store: MemoryStore, +) -> None: + capabilities = replace(DEFAULT_CAPABILITIES, max_rows_per_op=1) + writer = Endpoint(MemoryDriver(store, capabilities=capabilities), ["X", "Y", "Z"]) + writer.driver.fail("put", RuntimeError("lost")) # X's put + attempt = writer.backend.publish(writer.view("X", "Y", "Z"), None) + outcome = writer.wait_outcome(attempt) + assert masks(outcome) == [[False, False, False]] + assert outcome.fault == "group 0 row 0: the store driver raised " + repr(RuntimeError("lost")) + assert store_calls(writer) == ["put"] # Y and Z never reached the driver + with pytest.raises(SubmissionRejected, match="broke its protocol"): + writer.backend.publish(writer.view("X"), None) + + +@keeps_backend_threads +def test_a_violation_told_during_the_hand_over_stops_the_rest_of_it(store: MemoryStore) -> None: + capabilities = replace(DEFAULT_CAPABILITIES, max_rows_per_op=1) + writer = Endpoint(MemoryDriver(store, capabilities=capabilities), ["A", "B", "C"]) + for name in "ABC": + writer.write(name, name.encode() * SLOT_BYTES) + # A's put runs inside put and calls on_change there; that query raises. + writer.driver.fail("access_ended", RuntimeError("flaky")) + attempt = writer.backend.publish(writer.view("A", "B", "C"), None) + outcome = writer.wait_outcome(attempt) + assert store_calls(writer) == ["put"] # B and C never reached the driver + assert masks(outcome) == [[False, False, False]] + assert outcome.fault == "group 0 row 0: a store driver query raised " + repr( + RuntimeError("flaky") + ) + with pytest.raises(SubmissionRejected, match="broke its protocol"): + writer.backend.publish(writer.view("A"), None) + + +class BreaksItsProtocolDuringAPut(MemoryDriver): + """Controlled. A put of ``trigger_key`` first tells of the earlier operation ``victim`` on + another thread, with a poll one result short: a violation recorded while a submission is + being handed over.""" + + victim: MemoryOp | None = None + trigger_key: str | None = None + short = False + + def poll(self, op): + results = super().poll(op) + return results[:-1] if self.short and op is self.victim else results + + def put(self, rows, *, on_change=None): + op = super().put(rows, on_change=on_change) + if self.trigger_key in rows.keys(): + self.short = True + teller = threading.Thread(target=self.victim.on_change, args=(self.victim,)) + teller.start() + teller.join() + return op + + +@keeps_backend_threads +def test_calls_cut_off_by_another_attempts_violation_end_their_access(store: MemoryStore) -> None: + driver = BreaksItsProtocolDuringAPut(store, capabilities=ONE_ROW_PER_CALL, controlled=True) + writer = Endpoint(driver, ["X", "A", "B", "C"]) + for name in "XABC": + writer.write(name, name.encode() * SLOT_BYTES) + first = writer.backend.publish(writer.view("X"), None) + driver.victim, driver.trigger_key = writer.op("X"), store_key("A") + second = writer.backend.publish(writer.view("A", "B", "C"), None) + assert store_calls(writer) == ["put", "put"] # X, then A; B and C never reached the driver + op_a = writer.op("A") + driver.complete(op_a) + driver.end_access(op_a) + outcome = writer.wait_outcome(second) + assert masks(outcome) == [[True, False, False]] + assert "group 0 row 1: not handed over: the store driver broke its protocol" in outcome.fault + with pytest.raises(SubmissionRejected, match="broke its protocol"): + writer.backend.publish(writer.view("B"), None) + # B and C touched nothing, so the second attempt quiesces; only the first stays held. + writer.wait_quiesced(second) + assert writer.backend.shutdown(0) == frozenset({first.label, "region:0"}) + + +@keeps_backend_threads +def test_a_log_line_that_blocks_never_holds_up_a_submission( + store: MemoryStore, monkeypatch: pytest.MonkeyPatch +) -> None: + logging_now, let_the_log_go = threading.Event(), threading.Event() + + def blocking_warning_once(*args, **kwargs) -> None: # as a full log pipe would block + logging_now.set() + let_the_log_go.wait(10 * WAIT_S) # outlasts the submitter's wait below + + monkeypatch.setattr(backend_module.logger, "warning_once", blocking_warning_once) + driver = MemoryDriver(store, controlled=True) + endpoint = Endpoint(driver, ["A", "B"]) + endpoint.write("A", b"aaaaaaaa") + endpoint.write("B", b"bbbbbbbb") + endpoint.backend.publish(endpoint.view("A"), None) + op = endpoint.op("A") + # Its on_change records the call as unprovable, then logs that. + marker = threading.Thread(target=driver.mark_unprovable, args=(op,), daemon=True) + marker.start() + try: + assert logging_now.wait(WAIT_S) + published: list[object] = [] + submitter = threading.Thread( + target=lambda: published.append(endpoint.backend.publish(endpoint.view("B"), None)), + daemon=True, + ) + submitter.start() + submitter.join(WAIT_S) + assert published, "a submission waited for a log line" + finally: + let_the_log_go.set() + marker.join(WAIT_S) + + +@keeps_backend_threads +def test_a_protocol_violation_is_logged_once_and_never_under_the_lock( + store: MemoryStore, monkeypatch: pytest.MonkeyPatch +) -> None: + logged: list[str] = [] + logging_now, let_the_log_go = threading.Event(), threading.Event() + + def blocking_warning(message: str, *args, **kwargs) -> None: # as a full log pipe would block + logged.append(message) + logging_now.set() + let_the_log_go.wait(10 * WAIT_S) # outlasts the waits below + + monkeypatch.setattr(backend_module.logger, "warning", blocking_warning) + writer = Endpoint(MemoryDriver(store), ["A", "B"]) + writer.driver.fail("put", DriverRejected("busy")) + writer.wait_outcome(writer.backend.publish(writer.view("A"), None)) + assert logged == [] # a refusal breaks no protocol + writer.driver.fail("put", RuntimeError("bug")) + violator = threading.Thread( + target=lambda: writer.backend.publish(writer.view("B"), None), daemon=True + ) + violator.start() + try: + assert logging_now.wait(WAIT_S) + stats_read: list[object] = [] + reader = threading.Thread( + target=lambda: stats_read.append(writer.backend.stats()), daemon=True + ) + reader.start() + reader.join(WAIT_S) + assert stats_read, "the backend's lock was held while a violation was logged" + finally: + let_the_log_go.set() + violator.join(WAIT_S) + assert len(logged) == 1 and "broke its protocol" in logged[0] + + +def test_a_wrong_row_count_after_the_access_end_was_recorded_leaves_the_call_provable( + store: MemoryStore, +) -> None: + commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = ShortPollDriver(store, controlled=True) + reader = Endpoint(driver, ["A", "B"]) + attempt = reader.backend.fetch(reader.view("A", "B"), None) + op = reader.op("A") + driver.transfer(op) + driver.end_access(op) + wait_until(lambda: reader.backend.stats()["calls_awaiting_access_end"] == 0, "the access end") + driver.short = True # the next poll, now that the access end is recorded + driver.publish_results(op) + outcome = reader.wait_outcome(attempt) + assert "gave 1 results for 2 rows" in outcome.fault + assert attempt.quiesced() + assert reader.backend.stats()["calls_unprovable"] == 0 + with pytest.raises(SubmissionRejected, match="broke its protocol"): # a violation all the same + reader.backend.fetch(reader.view("A"), None) + reader.close() + + +@keeps_backend_threads +@pytest.mark.parametrize("failure", ["access_ended_raises", "poll_answers_no_sequence"]) +def test_a_query_that_fails_holds_the_call_and_refuses_later_work( + store: MemoryStore, failure: str +) -> None: + writer = Endpoint(MemoryDriver(store), ["A", "B"]) + writer.write("A", b"aaaaaaaa") + if failure == "access_ended_raises": + writer.driver.fail("access_ended", RuntimeError("flaky")) + else: + writer.driver.answer("poll", None) # the first query of the first operation + attempt = writer.backend.publish(writer.view("A"), None) + outcome = writer.wait_outcome(attempt) + assert masks(outcome) == [[False]] and "a store driver query raised" in outcome.fault + assert not attempt.quiesced() # no later answer of the driver proves the access end + assert writer.backend.stats()["calls_unprovable"] == 1 + with pytest.raises(SubmissionRejected, match="broke its protocol: a store driver query raised"): + writer.backend.publish(writer.view("B"), None) + assert writer.backend.shutdown(0) == frozenset({attempt.label, "region:0"}) + + +def test_a_raising_wake_does_not_stop_the_backend(store: MemoryStore) -> None: + def broken_wake() -> None: + raise RuntimeError("holder gone") + + writer = Endpoint(MemoryDriver(store), ["A", "B"], on_wake=broken_wake) + writer.write("A", b"aaaaaaaa") + writer.write("B", b"bbbbbbbb") + first = writer.backend.publish(writer.view("A"), None) + wait_until(lambda: first.poll() is not None and first.quiesced(), "A published") + second = writer.backend.publish(writer.view("B"), None) + wait_until(lambda: second.poll() is not None and second.quiesced(), "B published") + assert masks(first.poll()) == masks(second.poll()) == [[True]] + writer.close() + + +# Submission. + + +def _view_with(endpoint: Endpoint, **changes) -> CpuView: + rows = endpoint.rows(["A"]) + fields = { + "layer_group_index": 0, + "object_keys": rows.object_keys, + "addresses": rows.addresses, + "region_index": 0, + } + fields.update(changes) + return CpuView([Rows(**fields)]) + + +INVALID_SUBMISSIONS: dict[str, Callable[[Endpoint], tuple[CpuView, object]]] = { + "hint": lambda e: (e.view("A"), "peer-0"), + "no_keys": lambda e: (_view_with(e, object_keys=None, addresses=None, region_index=None), None), + "key_dtype": lambda e: (_view_with(e, object_keys=np.zeros((1, 4), dtype=np.int32)), None), + "key_width": lambda e: (_view_with(e, object_keys=np.zeros((1, 0), dtype=np.uint8)), None), + "address_dtype": lambda e: (_view_with(e, addresses=e.region.slots(0).astype(np.uint64)), None), + "len_disagrees_with_arrays": lambda e: (_view_with(e, rows=2), None), + "past_the_region": lambda e: (_view_with(e, addresses=e.region.slots(2)), None), + "before_the_region": lambda e: (_view_with(e, addresses=e.region.slots(-1)), None), + "misaligned": lambda e: (_view_with(e, addresses=e.region.slots(0) + 1), None), + "unknown_region_index": lambda e: (_view_with(e, region_index=1), None), + "negative_region_index": lambda e: (_view_with(e, region_index=-1), None), + "same_slot_in_a_group": lambda e: ( + CpuView([Rows(0, np.stack([object_key("A"), object_key("B")]), e.region.slots(0, 0), 0)]), + None, + ), + "same_slot_across_groups": lambda e: (CpuView([e.rows(["A"], 0), e.rows(["A"], 1)]), None), +} + + +def test_object_keys_in_any_memory_order_are_taken(store: MemoryStore) -> None: + commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + reader = Endpoint(MemoryDriver(store), ["A", "B"]) + rows = reader.rows(["A", "B"]) + view = CpuView([Rows(0, np.asfortranarray(rows.object_keys), rows.addresses, 0)]) + attempt = reader.backend.fetch(view, None) + assert masks(reader.wait_outcome(attempt)) == [[True, True]] + reader.wait_quiesced(attempt) + reader.close() + + +@pytest.mark.parametrize("case", INVALID_SUBMISSIONS) +def test_invalid_submissions_are_rejected_before_the_driver(store: MemoryStore, case: str) -> None: + reader = Endpoint(MemoryDriver(store), ["A", "B"]) + view, hint = INVALID_SUBMISSIONS[case](reader) + with pytest.raises(InvalidSubmission): + reader.backend.fetch(view, hint) + assert store_calls(reader) == [] + assert reader.backend.stats()["attempts_active"] == 0 + reader.close() + + +def test_a_fetch_hint_reaches_every_get_of_a_driver_that_reads_hints(store: MemoryStore) -> None: + commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb", "C": b"cccccccc"}) + capabilities = replace(DEFAULT_CAPABILITIES, hints=True, max_rows_per_op=2) + driver = MemoryDriver(store, capabilities=capabilities) + reader = Endpoint(driver, ["A", "B", "C"]) + source = ("peer", 0) + attempt = reader.backend.fetch(reader.view("A", "B", "C"), source) + assert masks(reader.wait_outcome(attempt)) == [[True, True, True]] + assert len(reader.driver.calls_to("get")) == 2 + assert driver.hints_seen == [source, source] and driver.hints_seen[0] is source + with pytest.raises(InvalidSubmission, match="None destination"): # a put takes no hint + reader.backend.publish(reader.view("A"), source) + reader.wait_quiesced(attempt) + unhinted = reader.backend.fetch(reader.view("A"), None) + assert masks(reader.wait_outcome(unhinted)) == [[True]] + assert driver.hints_seen == [source, source] # a None hint is never passed on + reader.wait_quiesced(unhinted) + reader.close() + + +def test_a_publish_may_read_one_slot_twice(store: MemoryStore) -> None: + writer = Endpoint(MemoryDriver(store), ["A", "B"]) + writer.write("A", b"aaaaaaaa") + keys = np.stack([object_key("A"), object_key("B")]) + view = CpuView([Rows(0, keys, writer.region.slots(0, 0), 0)]) + attempt = writer.backend.publish(view, None) + assert masks(writer.wait_outcome(attempt)) == [[True, True]] + assert store.committed_payload(store_key("B")) == b"aaaaaaaa" + writer.wait_quiesced(attempt) + writer.close() + + +def test_submissions_call_put_and_get_and_queries_run_on_the_driver_threads( + store: MemoryStore, +) -> None: + commit_objects(store, {"A": b"aaaaaaaa"}) + driver = MemoryDriver(store, capabilities=ONE_ROW_PER_CALL, background=True) + endpoint = Endpoint(driver, ["A", "B"]) + endpoint.write("B", b"bbbbbbbb") + fetched = endpoint.backend.fetch(endpoint.view("A"), None) + offered = endpoint.backend.publish(endpoint.view("B"), None) + for attempt in (fetched, offered): + assert masks(endpoint.wait_outcome(attempt)) == [[True]] + endpoint.wait_quiesced(attempt) + submitter = threading.current_thread().name + handed_over = [c for c in endpoint.driver.calls if c.method in ("put", "get")] + queries = [c for c in endpoint.driver.calls if c.method in ("poll", "access_ended")] + assert len(handed_over) == 2 and all(c.thread == submitter for c in handed_over) + assert queries and all(c.thread == "memory-driver-op" for c in queries) + endpoint.close() + + +def test_submissions_return_while_the_driver_waits_on_the_store() -> None: + store = KeyGatedStore(gets=True) + commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}) + driver = one_row_per_call(store, background=True) + endpoint = Endpoint(driver, ["A", "B", "C", "D"]) + endpoint.write("C", b"cccccccc") + blocked = endpoint.backend.fetch(endpoint.view("A", "B"), None) + attempts = [blocked] + try: + store.wait_for_calls_at_gate(2) + for submit in ( + lambda: attempts.append(endpoint.backend.fetch(endpoint.view("D"), None)), + lambda: attempts.append(endpoint.backend.publish(endpoint.view("C"), None)), + blocked.cancel, + ): + started = time.monotonic() + submit() + assert time.monotonic() - started < PROMPT_S + # Both of the cancelled fetch's calls are still inside the store. + assert blocked.quiesced() is False + assert blocked.label in endpoint.backend.shutdown(0) + finally: + store.gate.set() + for attempt in attempts: + endpoint.wait_quiesced(attempt) + endpoint.close() + + +# Shutdown. + + +def test_shutdown_forces_fault_outcomes_and_waits_out_its_budget(store: MemoryStore) -> None: + commit_objects(store, {"B": b"short"}) + driver = one_row_per_call(store, controlled=True) + driver.set_cancel_effective(False) # the forwarded cancels end no access + reader = Endpoint(driver, ["A", "B"]) + attempt = reader.backend.fetch(reader.view("A", "B"), None) + op_a, op_b = reader.op("A"), reader.op("B") + driver.transfer(op_a) # running: the drain a shutdown starts cannot cancel it + driver.complete(op_b) # B's fault (the wrong size) lands; A has no result, neither access ends + assert reader.backend.stats()["calls_awaiting_results"] == 1 # recorded as complete told of it + started = time.monotonic() + labels = reader.backend.shutdown(0.2) + assert time.monotonic() - started >= 0.19 # both calls may still access their slots + outcome = attempt.poll() + assert outcome is not None and SHUTDOWN_FAULT in outcome.fault + assert "first fault: group 0 row 1: object is 5 bytes" in outcome.fault + assert masks(outcome) == [[False, False]] + assert labels == frozenset({attempt.label, "region:0"}) + with pytest.raises(RuntimeError, match="region:0"): + reader.binding.release_hold() + assert not reader.hold.released + reader.backend.shutdown(0) + assert attempt.poll() is outcome + driver.publish_results(op_a) # a late result + driver.end_access(op_a) + driver.end_access(op_b) + reader.wait_quiesced(attempt) + assert attempt.poll() is outcome + reader.close() + + +@keeps_backend_threads +def test_shutdown_returns_once_every_open_attempt_has_an_unprovable_call( + store: MemoryStore, +) -> None: + commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb", "C": b"cccccccc"}) + driver = one_row_per_call(store, controlled=True) + driver.set_cancel_effective(False) # the forwarded cancels end no access + reader = Endpoint(driver, ["A", "B", "C"]) + never_quiesces = reader.backend.fetch(reader.view("A", "B"), None) + other = reader.backend.fetch(reader.view("C"), None) + op_a, op_b, op_c = reader.op("A"), reader.op("B"), reader.op("C") + for op in (op_a, op_b, op_c): + driver.transfer(op) # running: the drain a shutdown starts cannot cancel them + driver.mark_unprovable(op_a) # its results still pending + wait_until(lambda: reader.backend.stats()["calls_unprovable"] == 1, "the count") + # C's attempt can still quiesce, so a shutdown waits for it; A's unprovable call keeps its + # attempt from ever quiescing, whatever B does, so a shutdown does not wait for it. + labels: list[frozenset[str]] = [] + shutter = threading.Thread(target=lambda: labels.append(reader.backend.shutdown(WAIT_S))) + shutter.start() + time.sleep(PROMPT_S) + assert shutter.is_alive() + driver.publish_results(op_c) + driver.end_access(op_c) + shutter.join(PROMPT_S) + assert labels == [frozenset({never_quiesces.label, "region:0"})] + assert other.quiesced() and not op_b.access_ended + started = time.monotonic() + assert reader.backend.shutdown(WAIT_S) == frozenset({never_quiesces.label, "region:0"}) + assert time.monotonic() - started < PROMPT_S + # An unprovable call stays held whatever the driver says of it later, so the answer stands. + for op in (op_a, op_b): + driver.publish_results(op) + driver.end_access(op) + assert reader.backend.stats()["calls_awaiting_access_end"] == 1 # B's end, not A's + assert not never_quiesces.quiesced() + assert reader.backend.shutdown(0) == frozenset({never_quiesces.label, "region:0"}) + + +@keeps_backend_threads +def test_shutdown_returns_once_a_call_in_the_driver_turns_unprovable(store: MemoryStore) -> None: + driver = MemoryDriver(store, controlled=True) + driver.set_cancel_effective(False) # the forwarded cancel ends no access + writer = Endpoint(driver, ["A"]) + writer.write("A", b"aaaaaaaa") + attempt = writer.backend.publish(writer.view("A"), None) + op = writer.op("A") + driver.transfer(op) # running: the drain a shutdown starts cannot cancel it + labels: list[frozenset[str]] = [] + shutter = threading.Thread(target=lambda: labels.append(writer.backend.shutdown(WAIT_S))) + shutter.start() + time.sleep(PROMPT_S) + assert shutter.is_alive() # the put may still end its access + driver.mark_unprovable(op) + shutter.join(PROMPT_S) + assert labels == [frozenset({attempt.label, "region:0"})] + + +@pytest.mark.parametrize("stop", ["cancel", "shutdown"]) +def test_cancel_and_shutdown_ask_the_driver_to_cancel_the_open_calls( + store: MemoryStore, stop: str +) -> None: + commit_objects(store, {"A": b"aaaaaaaa", "B": b"bbbbbbbb", "C": b"cccccccc"}) + driver = one_row_per_call(store, controlled=True) + reader = Endpoint(driver, ["A", "B", "C"]) + attempt = reader.backend.fetch(reader.view("A", "B", "C"), None) + ops = [reader.op(name) for name in "ABC"] + if stop == "cancel": + driver.complete(ops[0]) + driver.end_access(ops[0]) + wait_until( + lambda: reader.backend.stats()["calls_awaiting_access_end"] == 2, "A's access end" + ) + attempt.cancel() + else: + assert attempt.label in reader.backend.shutdown(0) + reader.wait_quiesced(attempt) # the driver cancelled the rows it had not started + outcome = attempt.poll() + cancelled = {id(call.op) for call in reader.driver.calls_to("cancel")} + if stop == "cancel": # A's row is delivered, without a fault + assert outcome.fault is None and masks(outcome) == [[True, False, False]] + assert cancelled == {id(op) for op in ops[1:]} # only the calls the driver still held + else: + assert SHUTDOWN_FAULT in outcome.fault and masks(outcome) == [[False, False, False]] + assert cancelled == {id(op) for op in ops} + assert reader.read("B") == reader.read("C") == bytes(SLOT_BYTES) + reader.close() + + +@pytest.mark.parametrize("before", ["idle", "after_round_trip", "after_refused_put"]) +def test_clean_shutdown_returns_well_before_its_budget_and_rejects_later_submissions( + store: MemoryStore, before: str +) -> None: + endpoint = Endpoint(MemoryDriver(store), ["A"]) + endpoint.write("A", b"aaaaaaaa") + if before == "after_refused_put": + endpoint.driver.fail("put", DriverRejected("refused")) + if before != "idle": + attempt = endpoint.backend.publish(endpoint.view("A"), None) + endpoint.wait_quiesced(attempt) + calls = len(store_calls(endpoint)) + started = time.monotonic() + assert endpoint.backend.shutdown(5.0) == frozenset() + assert time.monotonic() - started < PROMPT_S + for submit in (endpoint.backend.fetch, endpoint.backend.publish): + with pytest.raises(SubmissionRejected): + submit(endpoint.view("A"), None) + assert len(store_calls(endpoint)) == calls + endpoint.close() + + +def test_shutdown_waits_out_an_unbounded_budget_without_raising() -> None: + store = KeyGatedStore(gets=True) + commit_objects(store, {"A": b"aaaaaaaa"}) + reader = Endpoint(MemoryDriver(store, background=True), ["A"]) + reader.backend.fetch(reader.view("A"), None) + labels: list[frozenset[str]] = [] + # A daemon thread, so a shutdown that never returns fails the wait below instead of the run. + shutter = threading.Thread( + target=lambda: labels.append(reader.backend.shutdown(math.inf)), daemon=True + ) + try: + store.wait_for_calls_at_gate(1) + shutter.start() + time.sleep(PROMPT_S) + assert shutter.is_alive() # waiting, since the call is still in the store + finally: + store.gate.set() + shutter.join(WAIT_S) + assert labels == [frozenset()] + reader.close() + + +def test_shutdown_never_calls_the_driver_and_unregisters_in_the_background( + store: MemoryStore, +) -> None: + endpoint = Endpoint(MemoryDriver(store), ["A"]) + revoke = endpoint.driver.gate("unregister") + try: + started = time.monotonic() + assert endpoint.backend.shutdown(0) == frozenset({"region:0"}) + assert time.monotonic() - started < PROMPT_S + endpoint.driver.wait_for(lambda: endpoint.driver.calls_to("unregister"), "a revocation") + started = time.monotonic() + with pytest.raises(RuntimeError, match="region:0"): + endpoint.binding.release_hold() # reads the registrations without waiting on it + assert time.monotonic() - started < PROMPT_S + finally: + revoke.set() + endpoint.close() + background = [c for c in endpoint.driver.calls if c.method in ("drain", "unregister")] + assert background and all(call.thread.endswith("-progress") for call in background) + + +class WakeProbe: + """On each wake, reads the backend from a helper thread as a lease holder would. A wake made + under the backend's lock would leave the helper blocked past its bound. Before ``attempt`` is + known, as during the submission that wakes, it reads only the backend.""" + + def __init__(self) -> None: + self.attempt = None + self.backend = None + self.seen: list[dict] = [] + + def __call__(self) -> None: + attempt, backend, seen = self.attempt, self.backend, {} + if backend is None: + return + + def read() -> None: + seen["stats"] = backend.stats() + if attempt is not None: + seen["outcome"] = attempt.poll() + seen["quiesced"] = attempt.quiesced() + attempt.cancel() + seen["finished"] = True + + helper = threading.Thread(target=read, daemon=True) + helper.start() + helper.join(WAIT_S) + self.seen.append(seen) + + +def test_an_outcome_that_lands_before_access_ends_wakes_the_holder(store: MemoryStore) -> None: + commit_objects(store, {"A": b"aaaaaaaa"}) + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A"]) + attempt = reader.backend.fetch(reader.view("A"), None) + op = reader.op("A") + woken = reader.wake.count + driver.complete(op) # results, with access still open + assert attempt.poll() is not None and not attempt.quiesced() + assert reader.wake.count > woken + driver.end_access(op) + reader.wait_quiesced(attempt) + reader.close() + + +def test_a_refused_call_wakes_the_holder(store: MemoryStore) -> None: + writer = Endpoint(MemoryDriver(store), ["A"]) + writer.driver.fail("put", DriverRejected("busy")) + attempt = writer.backend.publish(writer.view("A"), None) + assert attempt.poll() is not None and attempt.quiesced() # landed inside publish + assert writer.wake.count >= 1 + writer.close() + + +@pytest.mark.parametrize("settled", ["at_hand_over", "later"]) +def test_wake_runs_after_state_is_published_and_outside_the_lock_on_record( + store: MemoryStore, settled: str +) -> None: + commit_objects(store, {"A": b"aaaaaaaa"}) + probe = WakeProbe() + # An automatic driver settles the operation inside ``get`` and tells of it there; a controlled + # one settles it later, telling of each change as it comes. + driver = MemoryDriver(store, controlled=settled == "later") + reader = Endpoint(driver, ["A"], on_wake=probe) + probe.backend = reader.backend + probe.attempt = attempt = reader.backend.fetch(reader.view("A"), None) + if settled == "later": + op = reader.op("A") + driver.complete(op) + driver.end_access(op) + reader.wait_quiesced(attempt) + reader.wait_outcome(attempt) + assert probe.seen and all(seen.get("finished") for seen in probe.seen) + # A wake from inside ``get`` comes before fetch returned the attempt to the probe, so only a + # later wake is sure to read it. + if settled == "later": + assert any(seen["outcome"] is not None for seen in probe.seen) + assert any(seen["quiesced"] for seen in probe.seen) + reader.close() + + +def test_wake_runs_after_state_is_published_and_outside_the_lock_on_forced_shutdown() -> None: + store = KeyGatedStore(gets=True) + commit_objects(store, {"A": b"aaaaaaaa"}) + probe = WakeProbe() + reader = Endpoint(MemoryDriver(store, background=True), ["A"], on_wake=probe) + probe.backend = reader.backend + probe.attempt = reader.backend.fetch(reader.view("A"), None) + labels = [] + shutter = threading.Thread(target=lambda: labels.append(reader.backend.shutdown(WAIT_S))) + try: + store.wait_for_calls_at_gate(1) + shutter.start() + reader.wake.wait_for(lambda: reader.wake.count >= 1, "a wake for the forced outcome") + assert shutter.is_alive() # it cannot return while the call is in the store + (seen,) = probe.seen + assert seen.get("finished") and SHUTDOWN_FAULT in seen["outcome"].fault + finally: + store.gate.set() + shutter.join(WAIT_S) + assert labels == [frozenset()] + reader.close() + + +def test_the_progress_thread_never_keeps_the_process_alive(store: MemoryStore) -> None: + endpoint = Endpoint(MemoryDriver(store), ["A"]) + (progress,) = [thread for thread in endpoint.threads if thread.name.endswith("-progress")] + assert progress.daemon + endpoint.close() + + +def test_an_idle_backend_runs_no_progress_round( + store: MemoryStore, monkeypatch: pytest.MonkeyPatch +) -> None: + rounds: list[float] = [] + forward_cancels = BlobBackend._forward_cancels + + def counted(backend: BlobBackend) -> None: + rounds.append(time.monotonic()) + forward_cancels(backend) + + monkeypatch.setattr(BlobBackend, "_forward_cancels", counted) + endpoint = Endpoint(MemoryDriver(store), ["A"]) + time.sleep(PROMPT_S) + assert len(rounds) <= 1 # at most the round the thread starts with + endpoint.close() + + +# Registration. + + +HOST_ONLY = replace(DEFAULT_CAPABILITIES, device_memory=False) + + +@pytest.mark.parametrize("device_memory", [False, True], ids=["host_only", "device_memory"]) +def test_region_binding_registers_the_regions_with_slots_the_driver_can_access( + store: MemoryStore, device_memory: bool +) -> None: + driver = MemoryDriver(store, capabilities=replace(HOST_ONLY, device_memory=device_memory)) + regions = ( + CpuRegion(SLOT_BYTES, 0), + CpuRegion(SLOT_BYTES, 1), + CpuRegion(SLOT_BYTES, 1, memory_kind=MemoryKind.DEVICE), + ) + binding = RegionBinding(regions, driver, FakeHold()) + binding.register() + kinds = [registration.layout.memory_kind for registration in driver.live_registrations()] + if device_memory: + assert binding.live_labels() == frozenset({"region:1", "region:2"}) + assert kinds == [MemoryKind.HOST, MemoryKind.DEVICE] + else: + assert binding.live_labels() == frozenset({"region:1"}) and kinds == [MemoryKind.HOST] + assert binding.try_unregister_all() and driver.live_registrations() == [] + + +def test_a_binding_with_no_region_to_register_is_rejected_and_lets_the_hold_go( + store: MemoryStore, +) -> None: + driver, hold = MemoryDriver(store, capabilities=HOST_ONLY), FakeHold() + regions = (CpuRegion(SLOT_BYTES, 0), CpuRegion(SLOT_BYTES, 1, memory_kind=MemoryKind.DEVICE)) + with pytest.raises(ConstructionRejected, match="no region"): + RegionBinding(regions, driver, hold).register() + assert hold.released and driver.live_registrations() == [] + + +def test_a_region_whose_layout_is_refused_registers_nothing_and_lets_the_hold_go( + store: MemoryStore, +) -> None: + driver, hold = MemoryDriver(store), FakeHold() + refused = CpuRegion(SLOT_BYTES, 4) + refused.slot_bytes = 0 # slots of no size: no layout can describe them + with pytest.raises(ConstructionRejected, match="region:1"): + RegionBinding((CpuRegion(SLOT_BYTES, 4), refused), driver, hold).register() + assert hold.released and driver.live_registrations() == [] + + +@pytest.mark.parametrize( + "failure", ["register_rejected", "register_overlaps", "register_escaped", "rollback_failed"] +) +def test_a_failed_registration_releases_the_hold_only_when_nothing_escaped(failure: str) -> None: + regions = (CpuRegion(SLOT_BYTES, 2), CpuRegion(SLOT_BYTES, 2)) + driver = (RollbackFailsDriver if failure == "rollback_failed" else MemoryDriver)(MemoryStore()) + if failure == "register_overlaps": # the driver refuses the second region: DriverRejected + regions = (regions[0], regions[0]) + elif failure != "rollback_failed": + driver.fail_next_register("no memlock", escaped=failure == "register_escaped") + hold = FakeHold() + binding = RegionBinding(regions, driver, hold) + with pytest.raises(RuntimeError) as raised: + binding.register() + if failure in ("register_rejected", "register_overlaps"): # rolled back + assert isinstance(raised.value, ConstructionRejected) + assert hold.released and driver.live_registrations() == [] + assert binding.live_labels() == frozenset() + return + assert not isinstance(raised.value, ConstructionRejected) + assert not hold.released and driver.live_registrations() + with pytest.raises(RuntimeError, match="region:0"): + binding.release_hold() # a registration may be live: the hold stays + assert not hold.released + + +def test_rows_in_a_region_the_backend_did_not_register_are_rejected(store: MemoryStore) -> None: + device = CpuRegion(SLOT_BYTES, 1, memory_kind=MemoryKind.DEVICE) + reader = Endpoint(MemoryDriver(store, capabilities=HOST_ONLY), ["A"], others=[device]) + assert reader.binding.live_labels() == frozenset({"region:0"}) + view = CpuView([Rows(0, object_key("A")[None, :], device.slots(0), 1)]) + with pytest.raises(InvalidSubmission, match="region 1 is not registered"): + reader.backend.fetch(view, None) + with pytest.raises(ValueError, match="region 1 is not registered"): + reader.backend.driver_call_limit(1) + assert store_calls(reader) == [] + reader.close() + + +# Repeated object keys. + + +REPEATED_KEYS = { + # name: (the driver's max_rows_per_op, the keys the view names at slots 0, 1, 2, ..., rows in + # each call) + "no_limit": (None, "ABA", [2, 1]), + "limit_of_two": (2, "AABC", [1, 2, 1]), + "every_row_alike": (None, "AAA", [1, 1, 1]), +} + + +@pytest.mark.parametrize("direction", ["publish", "fetch"]) +@pytest.mark.parametrize("case", REPEATED_KEYS) +def test_rows_naming_one_key_at_two_addresses_never_share_a_call( + store: MemoryStore, case: str, direction: str +) -> None: + """A view may name one key at two addresses; a RowBatch carries no key twice, so the group is + cut there as well as at the driver's row limit, the rows keeping their order.""" + max_rows_per_op, keys, expected = REPEATED_KEYS[case] + slots = [f"S{i}" for i in range(len(keys))] # one slot per row, named apart from the keys + capabilities = replace(DEFAULT_CAPABILITIES, max_rows_per_op=max_rows_per_op) + endpoint = Endpoint(MemoryDriver(store, capabilities=capabilities), slots) + if max_rows_per_op is None: + assert endpoint.backend.driver_call_limit(0) is None # no limit declared; still cut + payload = {key: key.encode() * SLOT_BYTES for key in set(keys)} + view = CpuView( + [ + Rows( + 0, + np.stack([object_key(key) for key in keys]), + endpoint.region.slots(*range(len(keys))), + 0, + ) + ] + ) + if direction == "publish": + for slot, key in zip(slots, keys): + endpoint.write(slot, payload[key]) + attempt = endpoint.backend.publish(view, None) + else: + commit_objects(store, payload) + attempt = endpoint.backend.fetch(view, None) + outcome = endpoint.wait_outcome(attempt) + assert outcome.fault is None and masks(outcome) == [[True] * len(keys)] + method = "put" if direction == "publish" else "get" + assert [len(call.rows) for call in endpoint.driver.calls_to(method)] == expected + if direction == "publish": + assert {key: store.committed_payload(store_key(key)) for key in set(keys)} == payload + else: + assert [endpoint.read(slot) for slot in slots] == [payload[key] for key in keys] + endpoint.wait_quiesced(attempt) + endpoint.close() + + +# Regions registered together. + + +def registration_calls(endpoint: Endpoint) -> list[str]: + names = ("register", "unregister", "register_regions", "unregister_regions") + return [call.method for call in endpoint.driver.calls if call.method in names] + + +def test_a_together_driver_gets_every_region_in_one_batch_and_gives_them_back_in_one( + store: MemoryStore, +) -> None: + driver = TogetherDriver(store) + others = [CpuRegion(SLOT_BYTES, 0), CpuRegion(SLOT_BYTES, 2)] + endpoint = Endpoint(driver, ["A", "B"], others=others) + assert registration_calls(endpoint) == ["register_regions"] + assert endpoint.binding.live_labels() == frozenset({"region:0", "region:2"}) + (batch,) = driver.live_batches() + addresses = [registration.span.address for registration in batch.registrations] + assert addresses == [endpoint.region.address, others[1].address] # in region order + commit_objects(store, {"A": b"aaaaaaaa"}) + outcome = endpoint.wait_outcome(endpoint.backend.fetch(endpoint.view("A", "B"), None)) + assert masks(outcome) == [[True, False]] and endpoint.read("A") == b"aaaaaaaa" + endpoint.close() + assert registration_calls(endpoint) == ["register_regions", "unregister_regions"] + assert driver.live_batches() == [] and driver.live_registrations() == [] + + +@pytest.mark.parametrize( + "failure", ["rejected", "nothing_registered", "revocable", "revocation_fails", "no_handle"] +) +def test_a_failed_batch_registration_releases_the_hold_only_once_nothing_is_registered( + failure: str, +) -> None: + regions = (CpuRegion(SLOT_BYTES, 2), CpuRegion(SLOT_BYTES, 2)) + driver = TogetherDriver(MemoryStore()) + spy = TogetherDriverSpy(driver) + if failure == "rejected": + spy.fail("register_regions", DriverRejected("no memlock")) + elif failure == "no_handle": + spy.fail("register_regions", RuntimeError("the client crashed")) + else: + driver.fail_next_register_regions("no memlock", escaped=failure != "nothing_registered") + if failure == "revocation_fails": + driver.fail_next_unregister_regions("device busy") + hold = FakeHold() + binding = RegionBinding(regions, spy, hold) + with pytest.raises(RuntimeError) as raised: + binding.register() + revocations = len(spy.calls_to("unregister_regions")) + if failure in ("rejected", "nothing_registered", "revocable"): + assert isinstance(raised.value, ConstructionRejected) + assert hold.released and driver.live_registrations() == [] + assert binding.live_labels() == frozenset() + assert revocations == (1 if failure == "revocable" else 0) + return + # Something may be registered: the regions keep their owner and the hold stays. + assert not isinstance(raised.value, ConstructionRejected) + assert not hold.released + assert binding.live_labels() == frozenset({"region:0", "region:1"}) + with pytest.raises(RuntimeError, match="region:0"): + binding.release_hold() + if failure == "revocation_fails": # a handle: a later revocation succeeds and frees the hold + assert revocations == 1 and driver.live_registrations() + assert binding.try_unregister_all() and driver.live_registrations() == [] + binding.release_hold() + assert hold.released + else: # nothing to revoke by: the hold is kept for good + assert revocations == 0 and not binding.try_unregister_all() + assert not hold.released + + +def test_a_busy_region_keeps_the_whole_batch_registered(store: MemoryStore) -> None: + driver = TogetherDriver(store, controlled=True) + endpoint = Endpoint(driver, ["A"], others=[CpuRegion(SLOT_BYTES, 1)]) + endpoint.write("A", b"aaaaaaaa") + attempt = endpoint.backend.publish(endpoint.view("A"), None) + op = endpoint.op("A") + assert not endpoint.binding.try_unregister_all() # refused: the put may still read A + assert endpoint.binding.live_labels() == frozenset({"region:0", "region:1"}) + assert driver.live_batches() + driver.complete(op) + driver.end_access(op) + endpoint.wait_quiesced(attempt) + endpoint.close() + assert driver.live_batches() == [] + + +# A backend thread that cannot start. + + +@pytest.mark.parametrize("cleanup", ["succeeds", "fails"]) +def test_a_backend_whose_thread_cannot_start_leaves_the_binding_to_its_owner( + store: MemoryStore, cleanup: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """The progress thread cannot start: the constructor raises, no backend thread survives, and + the binding its owner registered before the build keeps the registrations and the hold for + that owner to revoke and release. A revocation that fails keeps both, as after any build.""" + driver = MemoryDriver(store) + region, hold = CpuRegion(SLOT_BYTES, 1), FakeHold() + binding = RegionBinding((region,), driver, hold) + binding.register() + if cleanup == "fails": + (handle,) = driver.live_registrations() + driver.fail_next_unregister(handle, "device busy") + + def start(thread: threading.Thread) -> None: + raise RuntimeError("can't start new thread") + + before = set(threading.enumerate()) + with monkeypatch.context() as patched: + patched.setattr(threading.Thread, "start", start) + with pytest.raises(RuntimeError, match="can't start new thread"): + BlobBackend(binding, lambda: None, BlobConfig()) + assert set(threading.enumerate()) == before + assert binding.live_labels() == frozenset({"region:0"}) and not hold.released + if cleanup == "fails": + assert not binding.try_unregister_all() and driver.live_registrations() + assert binding.live_labels() == frozenset({"region:0"}) + with pytest.raises(RuntimeError, match="region:0"): + binding.release_hold() + assert not hold.released + return + assert binding.try_unregister_all() and driver.live_registrations() == [] + binding.release_hold() + assert hold.released diff --git a/tests/unittest/_torch/disaggregation/blob/test_blob_example.py b/tests/unittest/_torch/disaggregation/blob/test_blob_example.py new file mode 100644 index 000000000000..5d3d6d15fd2c --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_blob_example.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The blob backend end to end on the CPU: two endpoints sharing objects through one store, a +lease holder's decisions and leases around the outcome and the quiescence of an attempt, the +region hold at shutdown, and two backends living in one process. Each test is also the smallest +working example of its part of the picture; ``README.md`` points here.""" + +from __future__ import annotations + +import threading +from dataclasses import replace + +import pytest +from cpu_fixtures import ( + SLOT_BYTES, + WAIT_S, + CpuRegion, + Endpoint, + FakeLeaseHolder, + commit_object, + masks, + store_key, + wait_until, +) + +from tensorrt_llm._torch.disaggregation.blob.store import StoreError +from tensorrt_llm._torch.disaggregation.testing.memory_store import ( + DEFAULT_CAPABILITIES, + MemoryDriver, + MemoryStore, +) + +pytestmark = pytest.mark.cpu_only + + +@pytest.fixture +def store() -> MemoryStore: + return MemoryStore() + + +def test_two_endpoints_share_objects_through_the_store(store: MemoryStore) -> None: + """A writer publishes two rows; a reader fetches them and a third that nobody published.""" + writer = Endpoint(MemoryDriver(store), ["A", "B"]) + reader = Endpoint(MemoryDriver(store), ["A", "B", "C"], fill=0x2D) + writer.write("A", b"aaaaaaaa") + writer.write("B", b"bbbbbbbb") + publish = writer.backend.publish(writer.view("A", "B"), None) + assert masks(writer.wait_outcome(publish)) == [[True, True]] + writer.wait_quiesced(publish) + assert writer.read("A") == b"aaaaaaaa" # a publish reads its source and never writes it + assert store.committed_payload(store_key("A")) == b"aaaaaaaa" + + fetch = reader.backend.fetch(reader.view("A", "B", "C"), None) + outcome = reader.wait_outcome(fetch) + assert masks(outcome) == [[True, True, False]] and outcome.fault is None + assert reader.read("A") == b"aaaaaaaa" and reader.read("B") == b"bbbbbbbb" + assert reader.read("C") == b"-" * SLOT_BYTES # a miss leaves the slot alone + reader.wait_quiesced(fetch) + + writer.write("A", b"xxxxxxxx") # the store keeps its own copy: the source may change after + again = reader.backend.fetch(reader.view("A"), None) + assert masks(reader.wait_outcome(again)) == [[True]] and reader.read("A") == b"aaaaaaaa" + reader.wait_quiesced(again) + writer.close() + reader.close() + + +def test_a_fault_with_a_late_write_keeps_the_lease_until_the_attempt_quiesced( + store: MemoryStore, +) -> None: + """The holder decides as soon as the outcome lands, but a faulted row's transfer may still + be writing: the lease waits for the evidence that it stopped.""" + for name, data in {"A": b"aaaaaaaa", "B": b"bbbbbbbb"}.items(): + commit_object(store, name, data) + one_row_per_call = replace(DEFAULT_CAPABILITIES, max_rows_per_op=1) + driver = MemoryDriver(store, controlled=True, capabilities=one_row_per_call) + reader = Endpoint(driver, ["A", "B"], fill=0x2D) + holder = FakeLeaseHolder(reader) + attempt = holder.fetch("A", "B") + op_a, op_b = reader.op("A"), reader.op("B") + driver.complete(op_a) + driver.end_access(op_a) + driver.fail(op_b, "link reset") # B's result is in; its access is not over + reader.wait_outcome(attempt) + holder.on_wake() + assert holder.delivered(attempt) == [[True, False]] + assert attempt.poll().fault == "group 0 row 1: link reset" + assert attempt.label in holder.leased # decided, and still leased + driver.write_late(op_b, 0, b"late....") # what the kept lease protects the slot from + assert reader.read("B") == b"late...." + driver.end_access(op_b) + reader.wait_quiesced(attempt) + holder.on_wake() + assert attempt.label not in holder.leased + assert holder.delivered(attempt) == [[True, False]] # the decision stood + reader.close() + + +def test_a_deadline_decides_from_nothing_and_a_late_outcome_changes_no_decision( + store: MemoryStore, +) -> None: + commit_object(store, "A", b"aaaaaaaa") + driver = MemoryDriver(store, controlled=True) + reader = Endpoint(driver, ["A"], fill=0x2D) + holder = FakeLeaseHolder(reader) + attempt = holder.fetch("A") + op = reader.op("A") + holder.decide(attempt) # the holder's deadline: nothing has concluded + assert holder.delivered(attempt) is None and attempt.label in holder.leased + driver.complete(op) + driver.end_access(op) + reader.wait_outcome(attempt) + reader.wait_quiesced(attempt) + holder.on_wake() + assert holder.delivered(attempt) is None # decided once; the row landed but does not count + assert reader.read("A") == b"aaaaaaaa" + assert attempt.label not in holder.leased # the lease still waited for the evidence + reader.close() + + +def test_the_hold_goes_back_only_once_every_region_is_revoked(store: MemoryStore) -> None: + """Two regions under one hold. Region 0's revocation fails once and its retry is held at a + gate, so the state in between is fixed: region 1 revoked, region 0 not. The manager keeps + the hold, and the binding, until the retry succeeded; nothing frees it on its own.""" + endpoint = Endpoint(MemoryDriver(store), ["A"], others=[CpuRegion(SLOT_BYTES, 2)]) + holder = FakeLeaseHolder(endpoint) + endpoint.driver.fail("unregister", StoreError("device busy"), times=1) # region 0, first try + retry = endpoint.driver.gate("unregister", after=2) # region 1 goes through; the retry waits + try: + labels = holder.shut_down(0) + assert labels and labels <= {"region:0", "region:1"} and not endpoint.hold.released + endpoint.driver.wait_for( + lambda: len(endpoint.driver.calls_to("unregister")) == 3, "the retry at the gate" + ) + assert endpoint.binding.live_labels() == frozenset({"region:0"}) + with pytest.raises(RuntimeError, match="region:0"): + endpoint.binding.release_hold() # one region still registered keeps the memory + assert not endpoint.hold.released + finally: + retry.set() + wait_until(lambda: not endpoint.binding.live_labels(), "the retried revocation") + assert not endpoint.hold.released # revoked, but the manager has not been told yet + assert holder.shut_down(0) == frozenset() and endpoint.hold.released + assert endpoint.hold.release_threads == [threading.get_ident()] + for thread in endpoint.threads: + thread.join(WAIT_S) + assert not any(thread.is_alive() for thread in endpoint.threads) + + +def test_two_backends_in_one_process_serve_and_close_independently(store: MemoryStore) -> None: + first = Endpoint(MemoryDriver(store), ["A", "B"], fill=0x2D) + second = Endpoint(MemoryDriver(store), ["A", "B"], fill=0x2D) + first.write("A", b"aaaaaaaa") + second.write("B", b"bbbbbbbb") + attempts = [ + first.backend.publish(first.view("A"), None), + second.backend.publish(second.view("B"), None), + ] + for endpoint, attempt in zip((first, second), attempts, strict=True): + assert masks(endpoint.wait_outcome(attempt)) == [[True]] + endpoint.wait_quiesced(attempt) + crossed = [ + first.backend.fetch(first.view("B"), None), # what the other one published + second.backend.fetch(second.view("A"), None), + ] + for endpoint, attempt in zip((first, second), crossed, strict=True): + assert masks(endpoint.wait_outcome(attempt)) == [[True]] + endpoint.wait_quiesced(attempt) + assert first.read("B") == b"bbbbbbbb" and second.read("A") == b"aaaaaaaa" + first.close() # the other backend, with its own driver and threads, is untouched + later = second.backend.fetch(second.view("B"), None) + assert masks(second.wait_outcome(later)) == [[True]] + second.wait_quiesced(later) + second.close() diff --git a/tests/unittest/_torch/disaggregation/blob/test_blob_format.py b/tests/unittest/_torch/disaggregation/blob/test_blob_format.py index e3a4ce336e14..622d6f7c5001 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_blob_format.py +++ b/tests/unittest/_torch/disaggregation/blob/test_blob_format.py @@ -4,14 +4,17 @@ from __future__ import annotations +import numpy as np 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.keys import FORMAT_VERSION, encode_key, key_stem from tensorrt_llm._torch.disaggregation.blob.store import ( DriverCapabilities, DriverRejected, RegionLayout, + RegionSpec, + RowBatch, RowIO, RowResult, RowStatus, @@ -78,6 +81,8 @@ def test_span_and_layout_validate_their_numbers() -> None: with pytest.raises(ValueError): Span(0, 0) assert Span(0x1000, 64).overlaps(Span(0x1020, 64)) + assert not Span(0x1040, 8).overlaps(Span(0x1000, 64)) # adjacent, so back to back + assert not Span(0x1000, 64).overlaps(Span(0x1040, 8)) assert Span(0x1010, 16).within(Span(0x1000, 64)) for slot_bytes, slot_count in ((0, 1), (8, 0), (-8, 1)): with pytest.raises(ValueError): @@ -87,6 +92,15 @@ def test_span_and_layout_validate_their_numbers() -> None: assert RegionLayout(8, 4, "kv", MemoryKind.DEVICE).memory_kind is MemoryKind.DEVICE +def test_a_region_spec_is_a_valid_span_whose_layout_fits() -> None: + spec = RegionSpec(0x1000, 32, RegionLayout(8, 4)) + assert spec.span == Span(0x1000, 32) + with pytest.raises(ValueError): + RegionSpec(0x1000, 0, RegionLayout(8, 4)) + with pytest.raises(ValueError): + RegionSpec(0x1000, 24, RegionLayout(8, 4)) # the layout needs 32 bytes + + 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] @@ -100,6 +114,59 @@ def test_a_row_is_a_key_and_disjoint_spans_in_the_order_given() -> None: RowIO("k", (Span(0x1000, 64), Span(0x1020, 64))) +def test_a_batch_of_slots_builds_nothing_until_read() -> None: + names = np.array([[0xAB, 0x01], [0xCD, 0x02], [0xEF, 0x03]], dtype=np.uint8) + slots = np.array([4096, 8192, 12288], dtype=np.int64) + region = Span(4096, 3 * 4096) + batch = RowBatch.of_slots(key_stem("p"), names, slots, 4096, region) + assert len(batch) == 3 and batch.spans_per_row == 1 and batch.row_bytes == 4096 + # Both columns are the caller's arrays: a change to them shows, as nothing was copied or built. + names[0] = [0x12, 0x34] + assert np.shares_memory(batch.addresses, slots) + assert not batch.addresses.flags.writeable # a driver cannot write through it + assert batch.addresses[:, 0].tolist() == [4096, 8192, 12288] + assert batch.span_bytes == (4096,) + assert batch.keys() == ["p/v1/1234", "p/v1/cd02", "p/v1/ef03"] + assert batch.extents == (region,) + assert batch[1] == RowIO("p/v1/cd02", (Span(8192, 4096),)) + assert list(batch) == [batch[0], batch[1], batch[2]] # read row by row, as a Sequence[RowIO] + + +def test_a_slice_of_a_batch_is_the_batch_of_those_rows() -> None: + names = np.arange(8, dtype=np.uint8).reshape(4, 2) + batch = RowBatch.of_slots("s/", names, np.arange(4, dtype=np.int64) * 8, 8, Span(0, 32)) + part = batch[1:3] + assert isinstance(part, RowBatch) and len(part) == 2 + assert part.keys() == batch.keys()[1:3] + assert part.addresses[:, 0].tolist() == [8, 16] and part.row_bytes == 8 + with pytest.raises(ValueError, match="contiguous"): + batch[::2] + + +def test_a_batch_of_rows_is_checked() -> None: + a = RowIO("a", (Span(0, 4),)) + b = RowIO("b", (Span(10, 4),)) + batch = RowBatch.of([a, b]) + assert batch.keys() == ["a", "b"] and list(batch) == [a, b] + assert not batch.addresses.flags.writeable # a driver cannot write through it + assert batch.span_bytes == (4,) and batch.row_bytes == 4 + assert batch.extents == (Span(0, 14),) # the smallest span holding every row + split = RowIO("m", (Span(8, 2), Span(0, 3))) + other = RowIO("n", (Span(20, 2), Span(4, 3))) + pair = RowBatch.of([split, other]) + assert pair.spans_per_row == 2 and pair.span_bytes == (2, 3) and pair.row_bytes == 5 + assert pair.extents == (Span(8, 14), Span(0, 7)) # one per span, in span order + assert pair[1] == other + with pytest.raises(ValueError, match="at least one row"): + RowBatch.of([]) + with pytest.raises(ValueError, match="appears twice"): + RowBatch.of([a, RowIO("a", (Span(8, 4),))]) + with pytest.raises(ValueError, match="same span sizes"): + RowBatch.of([a, RowIO("c", (Span(8, 6),))]) # rows of different sizes + with pytest.raises(ValueError, match="same span sizes"): + RowBatch.of([a, split]) # different numbers of spans + + 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] diff --git a/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py index 98e691553725..0d484071153f 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py +++ b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py @@ -1,7 +1,8 @@ # 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.""" +"""The in-memory driver: the shared conformance checks run several ways, its controlled mode, its +batch-registering variant, and proof that every check fails against a driver that breaks exactly +the rule it guards.""" from __future__ import annotations @@ -10,6 +11,7 @@ import time import weakref from collections.abc import Callable, Sequence +from dataclasses import replace import pytest @@ -18,11 +20,16 @@ DriverCapabilities, DriverRejected, Hint, + OnChange, OpHandle, ProbeStatus, + RegionHandle, RegionLayout, + RegionSpec, RegistrationError, + RowBatch, RowIO, + RowResult, RowStatus, Span, StoreDriver, @@ -33,20 +40,28 @@ from tensorrt_llm._torch.disaggregation.testing.broken_drivers import ( BROKEN_ASYNC_CONFORMANCE, BROKEN_CONFORMANCE, + BROKEN_TOGETHER_CONFORMANCE, ) from tensorrt_llm._torch.disaggregation.testing.driver_conformance import ( SLOT, UNRELEASED_ENDPOINTS, + UNRELEASED_MEMORY, Allocation, DriverConformance, Endpoint, Harness, Unblock, key_of, + specs_of, statuses, wait_until, ) -from tensorrt_llm._torch.disaggregation.testing.host_memory import GUARD, HostBuffer, MemoryWindow +from tensorrt_llm._torch.disaggregation.testing.host_memory import ( + GUARD, + HostBuffer, + Memory, + MemoryWindow, +) from tensorrt_llm._torch.disaggregation.testing.memory_store import ( DEFAULT_CAPABILITIES, MAX_PROBES_IN_FLIGHT, @@ -56,10 +71,15 @@ MemoryOp, MemoryStore, StoredObject, + TogetherMemoryDriver, ) from tensorrt_llm._torch.disaggregation.testing.object_trailer import encode_trailer -pytestmark = pytest.mark.cpu_only +# An exception that escapes a thread fails the test. +pytestmark = [ + pytest.mark.cpu_only, + pytest.mark.filterwarnings("error::pytest.PytestUnhandledThreadExceptionWarning"), +] PREFIX = "m" # the prefix the controlled-mode endpoints below use KEY_A, KEY_B, KEY_C = (key_of(label, PREFIX) for label in "ABC") @@ -99,11 +119,15 @@ def get(self, key: str) -> StoredObject | None: class CrashingStore(MemoryStore): - """Raises something that is not a ``StoreError`` from inside a commit, as a binding bug would.""" + """Raises something that is not a ``StoreError`` from inside every commit and read, as a + binding bug would.""" def commit(self, key: str, payload: bytes, trailer: bytes) -> Commit: raise RuntimeError("binding returned garbage") + def get(self, key: str) -> StoredObject | None: + raise RuntimeError("binding returned garbage") + class GatedStore(MemoryStore): """Holds reads, writes and queries of ``held`` keys (every key when None) while the gate is @@ -160,6 +184,10 @@ class TestMemoryDriverConformance(DriverConformance): driver_cls: type[MemoryDriver] = MemoryDriver capabilities = DEFAULT_CAPABILITIES background = False # True: put/get return once recorded and a thread executes them + # Automatic mode runs each operation inside put/get, which the protocol forbids; the checks + # that hold a put or a get inside the store skip it, and TestAsyncMemoryDriverConformance runs + # them. + returns_before_the_store_answers = False def _driver(self, store: MemoryStore | None = None) -> MemoryDriver: # An empty store is falsy (it has a length), so test for None explicitly. @@ -182,6 +210,11 @@ def make_blocked_get(self) -> tuple[StoreDriver, str, threading.Event | None, Un store.gate.clear() return self._driver(store), STAGED_KEY, store.waiting, store.gate.set + def make_blocked_put(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) @@ -193,6 +226,15 @@ def make_full_store(self) -> StoreDriver: def make_crash_after_record(self) -> StoreDriver: return self._driver(CrashingStore()) + def make_driver_leaving_operations_unstarted(self) -> StoreDriver: + return self.driver_cls(MemoryStore(), controlled=True, capabilities=self.capabilities) + + def make_results_before_access_end( + self, + ) -> tuple[StoreDriver, Callable[[OpHandle], None], Callable[[OpHandle], None]]: + driver = self.driver_cls(MemoryStore(), controlled=True, capabilities=self.capabilities) + return driver, driver.complete, driver.end_access + def make_failing_presence(self) -> tuple[StoreDriver, str]: store = MemoryStore() store.fail_next_presence("connection dropped") @@ -217,8 +259,8 @@ 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 get(self, rows: RowBatch, *, on_change=None) -> MemoryOp: # type: ignore[override] + return super().get(rows, on_change=on_change) def probe(self, keys: Sequence[str], budget_s: float) -> tuple[ProbeStatus, ...]: # type: ignore[override] return super().probe(keys, budget_s) @@ -232,7 +274,7 @@ class TestAsyncMemoryDriverConformance(TestMemoryDriverConformance): """The same driver returning handles before the store answers, as a real SDK does.""" background = True - handles_before_completion = True + returns_before_the_store_answers = True SHARED_STORE = MemoryStore() @@ -241,6 +283,8 @@ class TestAsyncMemoryDriverConformance(TestMemoryDriverConformance): class TestSharedStoreConformance(DriverConformance): """Fresh clients of one persistent store: every check finds the previous checks' objects.""" + returns_before_the_store_answers = False # the in-memory driver's automatic mode + def make_driver(self) -> StoreDriver: return MemoryDriver(SHARED_STORE) @@ -249,6 +293,7 @@ class TestPoolDriverConformance(DriverConformance): """A driver that registered its memory once, when built; endpoints are carved out of it.""" registers_on_demand = False + returns_before_the_store_answers = False # the in-memory driver's automatic mode POOL_BYTES = 1 << 16 def pool_driver(self) -> MemoryDriver: @@ -257,12 +302,15 @@ def pool_driver(self) -> MemoryDriver: 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.handle: RegionHandle = driver.register( + self.pool.address, self.pool.nbytes, self.pool_layout() ) self.next_offset = 0 return driver + def pool_layout(self) -> RegionLayout: + return RegionLayout(4, self.POOL_BYTES // 4, "pool") + 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 @@ -274,6 +322,137 @@ def shutdown(self, driver: StoreDriver) -> None: driver.close() +class TestTogetherDriverConformance(TestPoolDriverConformance): + """The pool registered as one batch through ``RegistersRegionsTogether``, plus the batch + scenarios the in-memory driver can stage.""" + + registers_together = True + driver_cls: type[TogetherMemoryDriver] = TogetherMemoryDriver + capabilities = DEFAULT_CAPABILITIES + + def pool_driver(self) -> TogetherMemoryDriver: + return self.driver_cls(MemoryStore(), capabilities=self.capabilities) + + def make_driver(self) -> StoreDriver: + driver = self.pool_driver() + self.pool = HostBuffer(self.POOL_BYTES) + spec = RegionSpec(self.pool.address, self.pool.nbytes, self.pool_layout()) + self.handle = driver.register_regions([spec]) + self.next_offset = 0 + return driver + + def shutdown(self, driver: StoreDriver) -> None: + assert isinstance(driver, TogetherMemoryDriver) + driver.unregister_regions(self.handle) + driver.close() + + def make_together_registration(self) -> tuple[StoreDriver, list[Memory]]: + return self.pool_driver(), [HostBuffer(2 * SLOT), HostBuffer(2 * SLOT)] + + def make_failing_together_registrations(self) -> list[tuple[StoreDriver, list[Memory]]]: + staged: list[tuple[StoreDriver, list[Memory]]] = [] + for escaped in (False, True): + driver = self.pool_driver() + driver.fail_next_register_regions("store refused", escaped=escaped) + staged.append((driver, [HostBuffer(2 * SLOT), HostBuffer(2 * SLOT)])) + return staged + + def make_busy_together_registration( + self, + ) -> tuple[StoreDriver, list[Memory], str, threading.Event, Unblock]: + store = GatedStore(held={STAGED_KEY}) + store.gate.clear() + driver = self.driver_cls(store, capabilities=self.capabilities) + memory: list[Memory] = [HostBuffer(2 * SLOT), HostBuffer(2 * SLOT)] + return driver, memory, STAGED_KEY, store.waiting, store.gate.set + + +class TestTogetherMemoryDriver: + """What only the batch-registering in-memory driver's controls can stage.""" + + def _driver(self, **kwargs: object) -> TogetherMemoryDriver: + return TogetherMemoryDriver(MemoryStore(), **kwargs) # type: ignore[arg-type] + + def test_single_region_calls_are_rejected(self) -> None: + driver = self._driver() + memory = HostBuffer(SLOT) + with pytest.raises(DriverRejected): + driver.register(memory.address, SLOT, RegionLayout(SLOT, 1)) + handle = driver.register_regions(specs_of([memory])) + with pytest.raises(DriverRejected): + driver.unregister(handle) + assert driver.live_batches() == [handle] + driver.unregister_regions(handle) + assert driver.drain(TIMEOUT) is True + driver.close() + + def test_a_batch_is_checked_whole_before_anything_is_registered(self) -> None: + driver = self._driver() + memory = HostBuffer(2 * SLOT) + inside = RegionSpec(memory.address + SLOT, SLOT, RegionLayout(SLOT, 1)) + with pytest.raises(DriverRejected): + driver.register_regions([*specs_of([memory]), inside]) # overlap within the batch + with pytest.raises(DriverRejected): + driver.register_regions([]) + assert driver.live_registrations() == [] and driver.live_batches() == [] + host_only = DriverCapabilities( + get=True, put=True, cancel=False, probe=False, max_rows_per_op=None + ) + driver = self._driver(capabilities=host_only) + device = RegionSpec(memory.address, SLOT, RegionLayout(SLOT, 1, "kv", MemoryKind.DEVICE)) + with pytest.raises(DriverRejected): + driver.register_regions([device]) + assert driver.live_registrations() == [] and driver.live_batches() == [] + + def test_a_busy_batch_is_not_revoked(self) -> None: + driver = self._driver(controlled=True) + memory = HostBuffer(SLOT) + handle = driver.register_regions(specs_of([memory])) + memory.write(b"a" * SLOT, 0) + op = driver.put(RowBatch.of([RowIO(KEY_A, (Span(memory.address, SLOT),))])) + with pytest.raises(DriverRejected): + driver.unregister_regions(handle) + driver.transfer(op) + driver.publish_results(op) + driver.end_access(op) + driver.unregister_regions(handle) + assert driver.drain(TIMEOUT) is True + driver.close() + + @pytest.mark.parametrize("escaped", [True, False]) + def test_a_failed_batch_hands_out_its_handle_only_when_something_may_be_registered( + self, escaped: bool + ) -> None: + driver = self._driver() + memory = HostBuffer(2 * SLOT) + driver.fail_next_register_regions("store refused", escaped=escaped) + with pytest.raises(RegistrationError) as caught: + driver.register_regions(specs_of([memory])) + if escaped: + assert caught.value.handle is not None + assert len(driver.live_registrations()) == 1 + driver.unregister_regions(caught.value.handle) + else: + assert caught.value.handle is None + assert driver.live_registrations() == [] + assert driver.drain(TIMEOUT) is True + driver.close() + + def test_a_failed_revocation_keeps_the_batch_live_for_a_retry(self) -> None: + driver = self._driver() + memory = HostBuffer(SLOT) + handle = driver.register_regions(specs_of([memory])) + driver.fail_next_unregister_regions("rpc timeout") + with pytest.raises(StoreError): + driver.unregister_regions(handle) + assert driver.live_batches() == [handle] + assert driver.drain(TIMEOUT) is True + with pytest.raises(DriverRejected): + driver.close() # a registration is still live + driver.unregister_regions(handle) + driver.close() + + class TestMemoryDriver: """Behaviour only the in-memory driver's controls or stores can stage.""" @@ -680,7 +859,7 @@ def test_memory_driver_rejects_overlapping_spans_within_one_get(self) -> None: RowIO(KEY_B, (Span(reader.span(0).address + 4, SLOT),)), # overlaps row A's span ] with pytest.raises(DriverRejected): - driver.get(rows) + driver.get(RowBatch.of(rows)) assert reader.read_row(0) == reader.read_row(1) == b"-" * SLOT assert statuses(driver, driver.get(reader.rows())) == [RowStatus.READY, RowStatus.READY] @@ -691,7 +870,7 @@ def test_a_single_span_driver_rejects_a_multi_span_row(self) -> None: 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)))]) + driver.put(RowBatch.of([RowIO(KEY_A, (writer.span(0), writer.span(1)))])) assert statuses(driver, driver.put(writer.rows())) == [RowStatus.READY, RowStatus.READY] # Submission and drain. @@ -743,6 +922,27 @@ def refuse(thread: threading.Thread) -> None: assert statuses(driver, later) == [RowStatus.READY, RowStatus.READY] self._shutdown(driver) + def test_drain_waits_for_the_last_on_change_of_every_operation(self) -> None: + """A worker counts as finished only once it told of its operation: a drain that returned + True earlier would let the owner close the driver while the callback still queries it.""" + driver = self.make(MemoryStore(), background=True) + reader = endpoint(driver, ["A"]) + entered, release = threading.Event(), threading.Event() + seen: list[tuple[list[RowStatus], bool]] = [] + + def look(op: OpHandle) -> None: + entered.set() + assert release.wait(TIMEOUT) + seen.append((statuses(driver, op), driver.access_ended(op))) + + driver.get(reader.rows(), on_change=look) + assert entered.wait(TIMEOUT) + assert driver.drain(0) is False # the worker is still telling of its operation + release.set() + assert driver.drain(TIMEOUT) is True + assert seen == [([RowStatus.MISS], True)] + 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"]) @@ -1087,7 +1287,7 @@ def _run_check(suite: object, check: str) -> None: def _conformance( driver_cls: type[MemoryDriver], caps: DriverCapabilities, - base: type[TestMemoryDriverConformance] = TestMemoryDriverConformance, + base: type[DriverConformance] = TestMemoryDriverConformance, ) -> DriverConformance: suite = type("Suite", (base,), {"driver_cls": driver_cls, "capabilities": caps, "timeout": 0.3}) return suite() # type: ignore[no-any-return] @@ -1112,6 +1312,23 @@ def test_each_async_conformance_check_passes_the_driver_and_fails_the_one_breaki _run_check(_conformance(MemoryDriver, capabilities, asynchronous), check) with pytest.raises(CHECK_FAILED): _run_check(_conformance(broken, capabilities, asynchronous), check) + # A broken driver may claim an access ended while its own thread still runs it. + wait_until( + lambda: not any(t.name == "memory-driver-op" for t in threading.enumerate()), + "the broken driver's threads to finish", + TIMEOUT, + ) + + +@pytest.mark.parametrize("rule", sorted(BROKEN_TOGETHER_CONFORMANCE)) +def test_each_together_conformance_check_passes_the_driver_and_fails_the_one_breaking_its_rule( + rule: str, +) -> None: + broken, capabilities, check = BROKEN_TOGETHER_CONFORMANCE[rule] + together = TestTogetherDriverConformance + _run_check(_conformance(TogetherMemoryDriver, capabilities, together), check) + with pytest.raises(CHECK_FAILED): + _run_check(_conformance(broken, capabilities, together), check) @pytest.mark.parametrize("rule", sorted(BROKEN_MEMORY)) @@ -1132,12 +1349,110 @@ def test_a_skip_on_the_broken_side_is_not_an_expected_failure() -> None: _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()] +class _CannotTellAbsenceFromFailure(TogetherMemoryDriver): + """A driver whose store cannot tell absence from failure: a get of a key the store does not + hold concludes FAULT and leaves its access unprovable, which the contract allows; every other + operation ends its access as usual. ``gets`` records the statuses of every get.""" + + gets: list[list[RowStatus]] = [] + + def __init__(self, store: MemoryStore, **options: object) -> None: + super().__init__(store, controlled=True, **options) + + def _get_row(self, io: RowIO) -> RowResult: + result = super()._get_row(io) + if result.status is RowStatus.MISS: + raise StoreError("the store cannot tell absence from failure") + return result + + def put(self, rows: RowBatch, *, on_change: OnChange | None = None) -> MemoryOp: + op = super().put(rows, on_change=on_change) + self.complete(op) + self.end_access(op) + return op + + def get( + self, rows: RowBatch, hint: Hint = None, *, on_change: OnChange | None = None + ) -> MemoryOp: + op = super().get(rows, hint, on_change=on_change) + self.complete(op) + self.gets.append(statuses(self, op)) + if RowStatus.FAULT in self.gets[-1]: + self.mark_unprovable(op) + else: + self.end_access(op) + return op + + +@pytest.mark.parametrize( + ("capabilities", "check", "region"), + [ + pytest.param( + replace(DEFAULT_CAPABILITIES, put=False, absent_is_miss=False), + "test_regions_registered_together_are_revoked_together_and_the_range_freed", + 0, + id="get-only probe", + ), + pytest.param( + replace(DEFAULT_CAPABILITIES, put=False, absent_is_miss=False), + "test_a_batch_is_not_revoked_while_any_of_its_regions_is_busy", + 0, + id="get-only probe of a busy batch", + ), + pytest.param( + replace(DEFAULT_CAPABILITIES, absent_is_miss=False), + "test_a_batch_is_not_revoked_while_any_of_its_regions_is_busy", + 1, + id="stalled get", + ), + ], +) +def test_a_row_left_unprovable_skips_the_batch_check_and_retains_the_memory( + capabilities: DriverCapabilities, check: str, region: int +) -> None: + """The batch checks move a row through each region to prove it registered, and the busy check + waits for its stalled get; a store that leaves such a row's access unprovable ends the check + as a skip naming the region, before any revocation, with the batch memory retained. A get-only + driver's probe is a get of a fresh key, which faults at once; with put the probes succeed and + the stalled get, of a key nobody holds, faults once released.""" + _CannotTellAbsenceFromFailure.gets = [] + suite = _conformance(_CannotTellAbsenceFromFailure, capabilities, TestTogetherDriverConformance) + retained = len(UNRELEASED_MEMORY) + with pytest.raises(pytest.skip.Exception, match=f"region {region}: .*unprovable"): + getattr(suite, check)() + assert len(UNRELEASED_MEMORY) == retained + 2 # both regions of the batch + gets = _CannotTellAbsenceFromFailure.gets + assert gets and all(got == [RowStatus.FAULT] for got in gets), gets + + +def test_every_check_has_a_broken_driver_and_each_driver_one_pairing() -> None: + pairings = [ + *BROKEN_CONFORMANCE.values(), + *BROKEN_ASYNC_CONFORMANCE.values(), + *BROKEN_TOGETHER_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_")} + # A check no broken driver fails is one nothing proves can fail. + 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_")} + + +class TestMemoryDriverOnChange: + @pytest.mark.parametrize("stop", ["cancel", "drain"]) + def test_a_cancel_or_drain_shows_before_on_change_is_called(self, stop: str) -> None: + driver = MemoryDriver(MemoryStore(), controlled=True) + reader = endpoint(driver, ["A"]) + seen: list[tuple[list[RowStatus], bool]] = [] + op = driver.get( + reader.rows(), + on_change=lambda op: seen.append((statuses(driver, op), driver.access_ended(op))), + ) + if stop == "cancel": + driver.cancel(op) + else: + assert driver.drain(TIMEOUT) is True + assert seen == [([RowStatus.CANCELLED], True)] diff --git a/tests/unittest/_torch/disaggregation/blob/test_together_driver.py b/tests/unittest/_torch/disaggregation/blob/test_together_driver.py new file mode 100644 index 000000000000..13798a6a5037 --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/test_together_driver.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The KVCR-shaped test driver: the shared conformance checks in batch mode, and what its shape +adds on top of the batch-registering in-memory driver.""" + +from __future__ import annotations + +import pytest +from test_memory_driver import TestTogetherDriverConformance as _BatchConformance +from together_driver import KVCR_LIKE_CAPABILITIES, TogetherDriver + +from tensorrt_llm._torch.disaggregation.blob.store import DriverRejected, RowBatch, RowIO, Span +from tensorrt_llm._torch.disaggregation.testing.driver_conformance import SLOT, specs_of +from tensorrt_llm._torch.disaggregation.testing.host_memory import HostBuffer +from tensorrt_llm._torch.disaggregation.testing.memory_store import MemoryStore + +pytestmark = pytest.mark.cpu_only + +TIMEOUT = 5.0 + + +class TestKvcrShapedDriverConformance(_BatchConformance): + """The suite in batch mode, with the capabilities the KVCR driver will declare.""" + + driver_cls = TogetherDriver + capabilities = KVCR_LIKE_CAPABILITIES + + +def test_one_batch_at_a_time_as_a_client_opened_with_its_memory() -> None: + driver = TogetherDriver(MemoryStore()) + first, second = HostBuffer(SLOT), HostBuffer(SLOT) + handle = driver.register_regions(specs_of([first])) + with pytest.raises(DriverRejected, match="one batch"): + driver.register_regions(specs_of([second])) + driver.unregister_regions(handle) + again = driver.register_regions(specs_of([second])) # the client was closed: a new one opens + driver.unregister_regions(again) + assert driver.drain(TIMEOUT) is True + driver.close() + + +def test_drain_leaves_the_registrations_to_unregister_regions() -> None: + """A backend drains first and revokes after; a drain that waited for the registrations + would wait on itself.""" + driver = TogetherDriver(MemoryStore()) + memory = HostBuffer(SLOT) + handle = driver.register_regions(specs_of([memory])) + assert driver.drain(TIMEOUT) is True + assert driver.live_batches() == [handle] + with pytest.raises(DriverRejected): + driver.close() # the batch is still live + driver.unregister_regions(handle) + driver.close() + + +def test_hints_are_taken_and_recorded() -> None: + driver = TogetherDriver(MemoryStore()) + memory = HostBuffer(SLOT) + handle = driver.register_regions(specs_of([memory])) + hint = ("peer-3", 7) + op = driver.get(RowBatch.of([RowIO("t/v1/00", (Span(memory.address, SLOT),))]), hint) + assert driver.hints_seen == [hint] + assert driver.access_ended(op) + driver.unregister_regions(handle) + assert driver.drain(TIMEOUT) is True + driver.close() diff --git a/tests/unittest/_torch/disaggregation/blob/together_driver.py b/tests/unittest/_torch/disaggregation/blob/together_driver.py new file mode 100644 index 000000000000..8577ca777f3b --- /dev/null +++ b/tests/unittest/_torch/disaggregation/blob/together_driver.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A store driver with the shape the KVCR driver will have, for the backend tests. + +KVCR opens its client with the memory it will use and closes the client to let the memory go, so +the driver registers regions together, one batch for as long as the client is open. Its gets take +a hint naming the peer to read from. It has no cancel and no probe. Its drain waits for its own +operations only, never for the registrations, which go when the batch is revoked; a drain that +waited for them would wait on the revocation that follows it. Over the in-memory store, the driver +keeps every control of ``TogetherMemoryDriver``. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from tensorrt_llm._torch.disaggregation.blob.store import ( + DriverCapabilities, + DriverRejected, + RegionSpec, +) +from tensorrt_llm._torch.disaggregation.testing.memory_store import ( + MemoryStore, + TogetherMemoryDriver, + _Batch, +) + +KVCR_LIKE_CAPABILITIES = DriverCapabilities( + get=True, + put=True, + cancel=False, + probe=False, + max_rows_per_op=None, + hints=True, + absent_is_miss=False, +) + + +class TogetherDriver(TogetherMemoryDriver): + """``TogetherMemoryDriver`` with KVCR's shape: one live batch at a time, hints taken and + recorded in ``hints_seen``, no cancel and no probe.""" + + def __init__(self, store: MemoryStore, **options: Any) -> None: + options.setdefault("capabilities", KVCR_LIKE_CAPABILITIES) + super().__init__(store, **options) + + def register_regions(self, specs: Sequence[RegionSpec]) -> _Batch: + with self._lock: + if self.live_batches(): + raise DriverRejected( + "the client is open with its memory already: one batch at a time" + ) + return super().register_regions(specs)