diff --git a/tensorrt_llm/_torch/disaggregation/blob/README.md b/tensorrt_llm/_torch/disaggregation/blob/README.md index 223f1a0bc9c9..67037b2e47ca 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/README.md +++ b/tensorrt_llm/_torch/disaggregation/blob/README.md @@ -66,7 +66,10 @@ The boundary is `store.py`; its docstrings are the specification. The rules that 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 + latter. For a put, `access_ended` is the caller's leave to overwrite the source: a driver + that copied the payload into memory of its own may answer True at once; one whose store + reads the source itself answers False until that read is done. + `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 @@ -85,6 +88,27 @@ The boundary is `store.py`; its docstrings are the specification. The rules that 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. +7. **The SDK's concurrency rules are the driver's to keep.** An SDK that promises no thread + safety is called from one thread at a time; one driver thread making every SDK call, + submissions and polls alike, is one way to keep that. `put` and `get` still only record and + queue, whatever thread runs the store's work. +8. **Callbacks only record.** `on_change` is called outside the driver's locks, as `OnChange` + says. A callback the SDK itself makes from inside one of its calls should only append to a + list the driver works through afterwards, in the order the SDK reported; running the + driver's logic inside it re-enters the SDK. +9. **Names a peer matches are computed alike on both sides.** When a remote store finds memory + by a name the driver chose, as a pool or region name, both sides must derive it from the same + logical mapping. A layout label plus an ordinal among alike regions holds only while both + sides register their regions in the same order; a local registration index never does. +10. **Evidence from sampling is newer than the first `drain`, and covers the store's own + work.** A `drain` that rests on a sampled signal, such as an in-flight count the store + publishes, accepts a sample taken after the first `drain` began and after the last local + operation ended, and only if it shows the store doing nothing for peers through this + driver's memory. One sample of nothing in flight is not enough on its own: the store must + also have stopped taking work that would undo it, which for a store serving peers means the + peers submit no more. The backend polls `drain(0.0)` every teardown round; a repeated call + with no budget must not reset the progress made so far, so that once every piece of work + ended and its evidence is in, a call with no budget observes True. ### The conformance suite @@ -103,6 +127,10 @@ class TestMyDriverConformance(DriverConformance): 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`. + A store whose pools each hold one row size overrides `allocate_rows(slot_bytes, rows, fill, + memory_kind)` instead, which tells it the row size an endpoint is laid out in. +- **Peers**: the suite stages no remote peer, so a store that moves rows between peers brings + its own cross-process fault cases, such as a peer stopped or gone while it writes. - **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`. diff --git a/tensorrt_llm/_torch/disaggregation/blob/store.py b/tensorrt_llm/_torch/disaggregation/blob/store.py index 3f5a50091a43..65150f2bb91d 100644 --- a/tensorrt_llm/_torch/disaggregation/blob/store.py +++ b/tensorrt_llm/_torch/disaggregation/blob/store.py @@ -73,7 +73,9 @@ class RowStatus(enum.Enum): READY = "ready" """For a get: the payload landed whole and intact, and the driver will not write the row again. For a put: a complete object is readable under the key by any reader now, whoever wrote it. A - store's "already exists" counts only once that object is committed.""" + store's "already exists" counts only once that object is committed. Now, not for good: a + store may evict the object later, and one that keeps it in a process loses it when that + process exits; a get after that reports MISS or FAULT.""" MISS = "miss" """For a get: the store reported that it holds no object under the key, and the row's memory @@ -498,8 +500,9 @@ def put(self, rows: RowBatch, *, on_change: OnChange | None = None) -> OpHandle: Raises: 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. + exceed ``max_rows_per_op``, an extent lies outside every live registration, or + the driver can serve no more, as once its store failed under it; nothing + escaped. """ ... @@ -572,6 +575,13 @@ def drain(self, budget_s: float) -> bool: 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. + + The budget is the most a call may wait, not a wait the driver can count on: the backend + calls ``drain(0.0)`` every round of its teardown and never waits inside it. A call with + no budget must not reset what the driver has established since its first ``drain``: once + every operation and every piece of the driver's own work ended, and the evidence of that + is in, a call with no budget observes True. A driver that could turn True only inside a + call with a budget never drains under the backend. """ ... diff --git a/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py b/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py index 78a66fcc4b19..ed76d172bdab 100644 --- a/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py +++ b/tensorrt_llm/_torch/disaggregation/testing/broken_drivers.py @@ -273,6 +273,24 @@ def drain(self, budget_s: float) -> bool: return self._quiet() +class DrainNeedsABudget(MemoryDriver): + """Turns True only inside a call with a budget, however idle it is: a driver that takes the + budget as its one chance to wait, or one whose evidence must be newer than each call, which + no evidence ever is. The backend only ever calls ``drain(0.0)``.""" + + def drain(self, budget_s: float) -> bool: + return super().drain(budget_s) and budget_s > 0 + + +class DrainWaitsForTheRunningGet(MemoryDriver): + """Waits for every operation to end whatever its budget says: a drain with no budget blocks + for as long as a get sits in the store. The check fails it from a thread of its own instead + of hanging behind it.""" + + def drain(self, budget_s: float) -> bool: + return super().drain(threading.TIMEOUT_MAX) + + class DrainForcesAnEnd(MemoryDriver): """Once its budget ran out, declares the access of every operation still running ended.""" @@ -482,6 +500,16 @@ def access_ended(self, op: OpHandle) -> bool: return recorded.access_ended or any(row.started for row in recorded.rows) +class EndsAccessBeforeTheSourceIsRead(MemoryDriver): + """Calls a put's access over as soon as it is recorded, while the store has yet to read the + source: a caller that trusts it overwrites what the store is about to take.""" + + def access_ended(self, op: OpHandle) -> bool: + with self._lock: + recorded = self._op(op) + return recorded.access_ended or recorded.kind == "put" + + 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.""" @@ -679,6 +707,16 @@ def _move(self, kind: Kind, io: RowIO) -> RowResult: DEFAULT_CAPABILITIES, "test_drain_is_one_way", ), + "drain turns true without a budget": ( + DrainNeedsABudget, + DEFAULT_CAPABILITIES, + "test_drain_polled_without_a_budget_turns_true_once_nothing_runs", + ), + "drain returns with no budget while a get is in the store": ( + DrainWaitsForTheRunningGet, + DEFAULT_CAPABILITIES, + "test_drain_polled_without_a_budget_turns_true_once_nothing_runs", + ), "close needs a true drain": ( ClosesAnyTime, DEFAULT_CAPABILITIES, @@ -804,6 +842,11 @@ def _move(self, kind: Kind, io: RowIO) -> RowResult: DEFAULT_CAPABILITIES, "test_a_drain_that_ends_an_operation_calls_its_on_change", ), + "a put's access stays open while its source is still to be read": ( + EndsAccessBeforeTheSourceIsRead, + DEFAULT_CAPABILITIES, + "test_a_put_keeps_its_access_open_while_the_store_still_reads_the_source", + ), } diff --git a/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py b/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py index 64beb2d1aaae..b649f855ad50 100644 --- a/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py +++ b/tensorrt_llm/_torch/disaggregation/testing/driver_conformance.py @@ -55,6 +55,7 @@ SLOT = 8 ROW_LIMIT_CHECKED = 64 # a row limit this high is taken on trust, so the check stays small +OFF_THREAD_NAME = "driver-conformance-call" # the threads a check runs a driver call on # Lender names are 54 bytes and start with a 16-byte scope digest shared by every name of a # deployment. A key built from a truncated name collides across these. @@ -208,6 +209,24 @@ def _access_settled(driver: StoreDriver, op: OpHandle, timeout: float) -> bool: return driver.access_ended(op) +def _refuses_stray( + h: Harness, direction: str, start: Callable[[], OpHandle], stray: Memory +) -> None: + """``start`` submits a row with a span in no registration; the driver must refuse it with + ``DriverRejected``, the one raise that certifies nothing escaped. One that takes the row, or + raises anything else, may touch the stray span for as long as the operation runs, and no + endpoint retains that memory, so it is kept alive before the check fails.""" + try: + start() + except DriverRejected: + return + except BaseException: + UNRELEASED_MEMORY.append(stray) + raise + UNRELEASED_MEMORY.append(stray) + pytest.fail(f"{direction} took a row with a span in no registration") + + 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.""" @@ -234,6 +253,9 @@ def __init__(self, suite: DriverConformance, driver: StoreDriver, timeout: float self.prefix = "t" + secrets.token_hex(4) # keys of one check never meet another's self.endpoints: list[Endpoint] = [] self.ops: list[OpHandle] = [] + # Threads whose driver call did not return in time; cleanup asks nothing more of a + # driver while one of them is still inside it. + self.hung: list[threading.Thread] = [] self.done = False # Creating things the harness will clean up. @@ -249,7 +271,7 @@ def endpoint( slot: int = SLOT, memory_kind: MemoryKind = MemoryKind.HOST, ) -> Endpoint: - memory, handle = self.suite.allocate(slot * len(labels), fill, memory_kind) + memory, handle = self.suite.allocate_rows(slot, len(labels), fill, memory_kind) endpoint = Endpoint( self.driver, labels, @@ -355,10 +377,18 @@ def cleanup(self) -> list[str]: Whenever cleanup stops short of shutting the driver down, every endpoint is retained for good: memory the driver may still touch is never freed. An operation whose end the driver declares unprovable is allowed by the contract and reported as no problem; one that - neither ends nor says so is. A cleanup that did not finish may be called again. + neither ends nor says so is. A driver call a check left hanging is given the timeout to + return once the check released what it waited on; if it still has not, the driver is + asked nothing more, since a second call would hang the run too. A cleanup that did not + finish may be called again. """ if self.done: return [] + for worker in self.hung: + worker.join(self.timeout) + if any(worker.is_alive() for worker in self.hung): + self._retain() + return ["a driver call never returned; the driver and its memory were left alone"] problems: list[str] = [] unprovable = False for op in self.ops: @@ -452,6 +482,14 @@ def allocate(self, nbytes: int, fill: int, memory_kind: MemoryKind) -> Allocatio """ return HostBuffer(nbytes, fill), None + def allocate_rows( + self, slot_bytes: int, rows: int, fill: int, memory_kind: MemoryKind + ) -> Allocation: + """Memory for one endpoint of ``rows`` rows of ``slot_bytes`` each: by default + ``allocate`` of their total. A driver whose pools each hold one row size overrides this + one instead of ``allocate``, so it can carve the endpoint from the pool of that size.""" + return self.allocate(slot_bytes * rows, fill, memory_kind) + def release(self, memory: Memory) -> None: """Gives back memory that ``allocate`` handed out with a handle; the suite never unregisters such memory itself.""" @@ -478,6 +516,15 @@ def make_blocked_put(self) -> tuple[StoreDriver, str, threading.Event | None, Un with an event set once a put is stalled, or ``None`` for no such signal.""" return None + def make_put_still_reading_the_source( + self, + ) -> tuple[StoreDriver, str, threading.Event | None, Unblock] | None: + """A driver whose store stalls every put of the returned key before it has read all of + the source, until the callable runs, with an event set once a put is stalled, or + ``None`` for no such signal. While the gate is closed the put can still touch the + caller's memory; a gate the store reaches only once the source is read does not qualify.""" + 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.""" @@ -722,11 +769,9 @@ def test_a_row_outside_every_registration_is_rejected_before_any_effect(self) -> h.needs(put=True, get=True) h.publish({"A": b"aaaaaaaa"}) stray = HostBuffer(SLOT, fill=0x2D) - with pytest.raises(DriverRejected): - h.get([RowIO(h.key("A"), (stray.span(),))]) + _refuses_stray(h, "get", lambda: h.get([RowIO(h.key("A"), (stray.span(),))]), stray) assert stray.read() == b"-" * SLOT and stray.canaries_intact() - with pytest.raises(DriverRejected): - h.put([RowIO(h.key("S"), (stray.span(),))]) + _refuses_stray(h, "put", lambda: h.put([RowIO(h.key("S"), (stray.span(),))]), stray) reader = h.endpoint(["S"], fill=0x2D) (result,) = h.wait_for_results(h.get(reader.rows())) h.assert_absent(result, reader, 0) @@ -775,10 +820,8 @@ def test_every_span_of_a_row_is_checked_against_the_registrations(self) -> None: 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]) + _refuses_stray(h, "put", lambda: h.put([row]), stray) + _refuses_stray(h, "get", lambda: h.get([row]), stray) assert stray.read() == b"-" * SLOT and stray.canaries_intact() assert inside.read_row(0) == b"aaaaaaaa" @@ -977,13 +1020,61 @@ def put() -> 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 + # Whether its access has ended is the driver's to say: one that staged the + # payload in memory of its own may truthfully have ended it already. finally: unblock() putter.join(h.timeout) (result,) = h.wait_for_results(seen["op"]) assert result.status is RowStatus.READY + def test_a_put_keeps_its_access_open_while_the_store_still_reads_the_source(self) -> None: + """``access_ended`` is the caller's leave to overwrite a put's source. While the store has + yet to read it the answer is False, whatever the driver has recorded; a driver that + copied the payload into memory of its own cannot stage this and skips. One that hands + out the handle only once the store answered cannot be asked.""" + if not self.returns_before_the_store_answers: + pytest.skip("this driver returns the handle of a put only once the store answered") + staged = self.make_put_still_reading_the_source() + if staged is None: + pytest.skip("this driver cannot stage a put with its source still to be read") + driver, key, reading, unblock = staged + with self.harness(driver) as h: + h.needs(put=True) + source = h.endpoint(["P"]) + source.fill_row(0, b"pppppppp") + handed: dict[str, OpHandle] = {} + handle_back = threading.Event() + failures: list[BaseException] = [] + + def submit() -> None: + try: + handed["op"] = h.put([RowIO(key, (source.span(0),))]) + except BaseException as error: # noqa: BLE001 - reported from the main thread + failures.append(error) + finally: + handle_back.set() + + putter = threading.Thread(target=submit, daemon=True) + putter.start() + try: + assert handle_back.wait(h.timeout), "put waited for the store to answer" + if failures: + raise AssertionError(f"put raised: {failures[0]!r}") from failures[0] + if reading is not None: + assert reading.wait(h.timeout), "the put never reached its source" + # With the gate closed the source cannot have been read through, wherever the + # put is on its way there. + assert not driver.access_ended(handed["op"]), ( + "access ended with the source still to be read" + ) + finally: + unblock() + putter.join(h.timeout) + (result,) = h.wait_for_results(handed["op"]) + assert result.status is RowStatus.READY + wait_until(lambda: driver.access_ended(handed["op"]), "access to end", h.timeout) + 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 @@ -1055,24 +1146,32 @@ def test_capabilities_do_not_change(self) -> None: assert h.driver.capabilities == before def test_a_raise_after_the_operation_is_recorded_becomes_a_row_fault(self) -> None: + """Each declared direction is checked on a driver of its own: a driver may serve no more + once its store raised under it and refuse the next call, as ``put`` allows, so the second + direction never meets the first one's driver. A get of an unpublished key still reaches + the store, so a get-only driver is checked through its get alone.""" driver = self.make_crash_after_record() if driver is None: pytest.skip("this driver cannot stage a store that raises mid-operation") - with self.harness(driver) as h: - # 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: + capabilities = driver.capabilities + if not (capabilities.put or capabilities.get): + pytest.skip("this driver declares neither put nor get") + directions = [ + d for d, declared in (("put", capabilities.put), ("get", capabilities.get)) if declared + ] + for index, direction in enumerate(directions): + if index: + driver = self.make_crash_after_record() + assert driver is not None, "the hook staged a driver once and then none" + with self.harness(driver) as h: + start: Callable[[Sequence[RowIO]], OpHandle] + if direction == "put": + writer = h.endpoint(["A"]) + writer.fill_row(0, b"aaaaaaaa") + start, rows = h.put, writer.rows() + else: + reader = h.endpoint(["A"], fill=0x2D) + start, rows = h.get, reader.rows() try: op = start(rows) except Exception as error: # noqa: BLE001 - the check is that nothing escapes @@ -1398,12 +1497,15 @@ def revoke_while_busy() -> None: 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" + # Proven only now: a driver that runs its operations one after another would queue + # these rows behind the stalled get. A region the refused revocation let go stays + # gone, so the proof is as good after the get returned. + for index, region in enumerate(memory[:-1]): # the idle regions stayed registered + if not _prove_registered(driver, region, index, self.timeout): + _skip_unprovable(index) if not _access_settled(driver, slow["op"], self.timeout): _skip_unprovable(len(memory) - 1) driver.unregister_regions(handle) @@ -1469,6 +1571,116 @@ def test_drain_is_one_way(self) -> None: assert h.driver.drain(self.timeout) is True assert writer.read_row(0) == b"aaaaaaaa" + def test_drain_polled_without_a_budget_turns_true_once_nothing_runs(self) -> None: + """The backend drains with no budget, every teardown round, and never gives the driver a + wait. Polled that way, ``drain`` answers False while a get runs in the store and forces + no end on it; once the access of that get ended and nothing runs, it turns True without + any call having had a budget. A suite that stages no blocked get, or none with a signal that + the get is in the store, checks the second half alone: without the signal, a drain may + find the row unstarted and cancel it, which ends its access as it should.""" + staged = self.make_blocked_get() + if staged is None: + with self.harness() as h: + self._drain_without_a_budget_after_a_put(h) + return + driver, blocked_key, blocked, unblock = staged + if blocked is None: # the staged driver still gets the check, and its cleanup + unblock() + with self.harness(driver) as h: + self._drain_without_a_budget_after_a_put(h) + return + with self.harness(driver) as h: + h.needs(get=True) + slow = h.endpoint(["S"], fill=0x2D) + started: dict[str, OpHandle] = {} + failures: list[BaseException] = [] + + def slow_get() -> None: + try: + started["op"] = h.get([RowIO(blocked_key, (slow.span(0),))]) + except BaseException as error: # noqa: BLE001 - reported from the main thread + failures.append(error) + + # Off this thread: a blocking driver's get returns only once the store lets it. + getter = threading.Thread(target=slow_get, daemon=True) + getter.start() + try: + assert blocked.wait(h.timeout), "the slow get never reached the store" + assert self._drain_without_a_budget(h) is False, ( + "drain answered True with a get in the store" + ) + if "op" in started: # the handle came back before the store answered + assert not h.driver.access_ended(started["op"]), "drain forced an access end" + finally: + unblock() + getter.join(h.timeout) + assert not getter.is_alive(), "the stalled get never returned" + if failures: + raise AssertionError(f"the get raised: {failures[0]!r}") from failures[0] + op = started["op"] + h.wait_for_results(op) + # A store may answer the released get with something that leaves its access + # unprovable; the contract allows that, and drain then rightly stays False. + if not _access_settled(h.driver, op, h.timeout): + pytest.skip("the store left the released get with its access unprovable") + self._poll_drain_without_a_budget(h) + + def _drain_without_a_budget_after_a_put(self, h: Harness) -> None: + """The second half of the check alone: one put, when the driver declares it, then a + drain polled with no budget.""" + if h.driver.capabilities.put: + h.publish({"A": b"aaaaaaaa"}) + self._poll_drain_without_a_budget(h) + + def _drain_without_a_budget(self, h: Harness) -> bool: + """One ``drain(0.0)``, off this thread: a drain that waits in spite of its budget, for a + get this thread holds in the store, fails the check instead of hanging it.""" + return bool(self._off_thread(h, lambda: h.driver.drain(0.0), "a drain with no budget")) + + def _poll_drain_without_a_budget(self, h: Harness) -> None: + """Polls ``drain(0.0)`` as the backend does until it answers True; fails at the timeout. + No call in between has a budget, so a driver that can turn True only inside a wait of + its own never does here. Off this thread too, so one call that never returns is a + failure rather than a hang.""" + + def poll() -> None: + wait_until( + lambda: h.driver.drain(0.0) is True, + "a drain with no budget to answer True", + h.timeout, + ) + + self._off_thread(h, poll, "the drain polling", h.timeout + 1.0) + + @staticmethod + def _off_thread( + h: Harness, call: Callable[[], object], what: str, within: float | None = None + ) -> object: + """Runs ``call`` on a thread of its own and returns what it returned, or raises what it + raised; fails once it has not returned within ``within`` seconds (the harness timeout + by default), so a call that waits on something this thread holds cannot hang the + check. The harness is told of a call that did not return, so its cleanup does not walk + into the same driver again.""" + outcome: list[tuple[bool, object]] = [] + + def run() -> None: + try: + outcome.append((True, call())) + except BaseException as error: # noqa: BLE001 - re-raised on the calling thread + outcome.append((False, error)) + + worker = threading.Thread(target=run, name=OFF_THREAD_NAME, daemon=True) + worker.start() + worker.join(h.timeout if within is None else within) + if worker.is_alive(): + h.hung.append(worker) + pytest.fail(f"{what} did not return") + returned, value = outcome[0] + if not returned: + assert isinstance(value, BaseException) + raise value + return value + def test_close_needs_a_true_drain_and_no_live_registrations(self) -> None: """For a driver whose memory was registered up front, the suite's shutdown hook does the unregistering; the order it has to keep is the same.""" diff --git a/tensorrt_llm/_torch/disaggregation/testing/memory_store.py b/tensorrt_llm/_torch/disaggregation/testing/memory_store.py index 6daf2c486429..89260f656c44 100644 --- a/tensorrt_llm/_torch/disaggregation/testing/memory_store.py +++ b/tensorrt_llm/_torch/disaggregation/testing/memory_store.py @@ -107,7 +107,7 @@ class MemoryStore: Subclass and override ``commit``, ``get`` or ``presence`` to make a test's store block or fail; they may raise ``StoreError``, which the driver turns into a row fault or re-raises from - ``probe``. + ``probe``. Override ``reading_source`` to hold a put before the caller's memory is read. """ def __init__(self, capacity: int | None = None) -> None: @@ -117,6 +117,10 @@ def __init__(self, capacity: int | None = None) -> None: self._presence_failure: str | None = None self.capacity = capacity + def reading_source(self, key: str) -> None: + """The put of ``key`` is about to read the caller's memory. A store that stalls here + keeps the source in use, as one that pulls the payload itself would.""" + def commit(self, key: str, payload: bytes, trailer: bytes) -> Commit: new = StoredObject(bytes(payload), bytes(trailer)) with self._lock: @@ -717,6 +721,7 @@ def _move(self, kind: Kind, io: RowIO) -> RowResult: return faulted(f"{type(error).__name__}: {error}") def _put_row(self, io: RowIO) -> RowResult: + self.store.reading_source(io.key) payload = _gather(io.spans) trailer = self._trailer_for(io, payload) outcome = self.store.commit(io.key, payload, trailer) diff --git a/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py index 0d484071153f..b3eb8cf3030c 100644 --- a/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py +++ b/tests/unittest/_torch/disaggregation/blob/test_memory_driver.py @@ -43,6 +43,7 @@ BROKEN_TOGETHER_CONFORMANCE, ) from tensorrt_llm._torch.disaggregation.testing.driver_conformance import ( + OFF_THREAD_NAME, SLOT, UNRELEASED_ENDPOINTS, UNRELEASED_MEMORY, @@ -131,11 +132,13 @@ def get(self, key: str) -> StoredObject | None: class GatedStore(MemoryStore): """Holds reads, writes and queries of ``held`` keys (every key when None) while the gate is - closed, and says when one is waiting.""" + closed, and says when one is waiting. With ``at_the_source`` a put is held before the + driver reads the caller's memory, not only before it commits.""" - def __init__(self, held: set[str] | None = None) -> None: + def __init__(self, held: set[str] | None = None, *, at_the_source: bool = False) -> None: super().__init__() self.held = held + self.at_the_source = at_the_source self.gate = threading.Event() self.gate.set() self.waiting = threading.Event() @@ -145,6 +148,10 @@ def _hold(self, key: str) -> None: self.waiting.set() self.gate.wait() + def reading_source(self, key: str) -> None: + if self.at_the_source: + self._hold(key) + def commit(self, key: str, payload: bytes, trailer: bytes) -> Commit: self._hold(key) return super().commit(key, payload, trailer) @@ -215,6 +222,13 @@ def make_blocked_put(self) -> tuple[StoreDriver, str, threading.Event | None, Un store.gate.clear() return self._driver(store), STAGED_KEY, store.waiting, store.gate.set + def make_put_still_reading_the_source( + self, + ) -> tuple[StoreDriver, str, threading.Event | None, Unblock]: + store = GatedStore(held={STAGED_KEY}, at_the_source=True) + 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) @@ -1441,6 +1455,184 @@ def test_every_check_has_a_broken_driver_and_each_driver_one_pairing() -> None: assert checks <= {name for name in dir(TestMemoryDriver) if name.startswith("test_")} +# The suite's scenario hooks. + + +def test_the_source_gate_holds_a_put_before_the_source_is_read() -> None: + """What the hook promises, for the in-memory driver: while the gate is closed nothing of the + source was read, so the store ends up with what the source holds once the gate opens.""" + suite = _conformance(MemoryDriver, DEFAULT_CAPABILITIES, TestAsyncMemoryDriverConformance) + staged = suite.make_put_still_reading_the_source() + assert staged is not None + driver, key, reading, unblock = staged + assert reading is not None + with suite.harness(driver) as h: + source = h.endpoint(["P"]) + source.fill_row(0, b"pppppppp") + try: + op = h.put([RowIO(key, (source.span(0),))]) + assert reading.wait(h.timeout), "the put never reached its source" + source.fill_row(0, b"qqqqqqqq") + finally: + unblock() # on every exit, or the worker behind the gate never returns + (result,) = h.wait_for_results(op) + assert result.status is RowStatus.READY + reader = h.endpoint(["R"]) + h.wait_for_results(h.get([RowIO(key, (reader.span(0),))])) + assert reader.read_row(0) == b"qqqqqqqq" + + +# The suite's memory hooks. + + +def test_a_suite_overriding_only_allocate_is_asked_for_the_total_as_before() -> None: + asked: list[tuple[int, int, MemoryKind]] = [] + + class Suite(TestMemoryDriverConformance): + def allocate(self, nbytes: int, fill: int, memory_kind: MemoryKind) -> Allocation: + asked.append((nbytes, fill, memory_kind)) + return super().allocate(nbytes, fill, memory_kind) + + with Suite().harness() as h: + h.endpoint(["A", "B"], slot=4, fill=0x2D) + assert asked == [(8, 0x2D, MemoryKind.HOST)] + + +def test_a_suite_overriding_allocate_rows_learns_the_row_size_and_count() -> None: + """Two endpoints of the same total, laid out differently: the hook tells them apart, which + a store whose pools each hold one row size needs to pick the pool.""" + asked: list[tuple[int, int, int, MemoryKind]] = [] + + class Suite(TestMemoryDriverConformance): + def allocate_rows( + self, slot_bytes: int, rows: int, fill: int, memory_kind: MemoryKind + ) -> Allocation: + asked.append((slot_bytes, rows, fill, memory_kind)) + return super().allocate_rows(slot_bytes, rows, fill, memory_kind) + + with Suite().harness() as h: + h.endpoint(["A", "B"], slot=4) + h.endpoint(["C"], slot=8) + assert asked == [(4, 2, 0, MemoryKind.HOST), (8, 1, 0, MemoryKind.HOST)] + + +def test_the_zero_budget_drain_check_closes_a_staged_driver_that_signals_nothing() -> None: + """A blocked-get hook without a signal still staged a driver; the check runs its second half + on that driver and closes it, rather than leaving it open beside a fresh one.""" + closed: list[MemoryDriver] = [] + staged: list[MemoryDriver] = [] + + class RecordsItsClose(MemoryDriver): + def close(self) -> None: + closed.append(self) + super().close() + + class Suite(TestMemoryDriverConformance): + driver_cls = RecordsItsClose + + def make_blocked_get(self) -> tuple[StoreDriver, str, threading.Event | None, Unblock]: + store = GatedStore(held={STAGED_KEY}) + store.gate.clear() + driver = self._driver(store) + staged.append(driver) + return driver, STAGED_KEY, None, store.gate.set + + Suite().test_drain_polled_without_a_budget_turns_true_once_nothing_runs() + assert closed == staged and len(staged) == 1 + + +class _TakesStraySpansAndNeverProvesItsEnd(MemoryDriver): + """Takes a row with a span in no registration, then can never prove that operation's access + ended: the worst a driver can do with memory the suite does not otherwise retain. Rows inside + the registrations move as usual.""" + + def _check_registered(self, rows: Sequence[RowIO]) -> None: + pass + + def _check_extents(self, rows: Sequence[RowIO]) -> None: + pass + + def _execute(self, op: MemoryOp) -> None: + if all(self._registered(span) for row in op.rows for span in row.io.spans): + super()._execute(op) + else: + self.mark_unprovable(op) + + +class _RaisesOnAStraySpan(_TakesStraySpansAndNeverProvesItsEnd): + """Lets a raise that is not ``DriverRejected`` out of a put or get of a row with a span in no + registration, after the operation was recorded: the memory may be in use all the same.""" + + def _execute(self, op: MemoryOp) -> None: + super()._execute(op) + if op.unprovable: + raise RuntimeError("the binding failed with the row in flight") + + +@pytest.mark.parametrize( + ("driver_cls", "outcome"), + [(_TakesStraySpansAndNeverProvesItsEnd, CHECK_FAILED), (_RaisesOnAStraySpan, RuntimeError)], +) +@pytest.mark.parametrize( + "check", + [ + "test_a_row_outside_every_registration_is_rejected_before_any_effect", + "test_every_span_of_a_row_is_checked_against_the_registrations", + ], +) +def test_a_stray_span_a_driver_took_is_retained_when_its_rejection_check_fails( + check: str, + driver_cls: type[MemoryDriver], + outcome: type[BaseException] | tuple[type[BaseException], ...], +) -> None: + """A driver that took the row fails the check; one that raised something other than + ``DriverRejected`` lets that raise through. Either way the stray memory the operation may + still touch is kept alive, as the endpoints of an unprovable operation are.""" + suite = _conformance(driver_cls, DEFAULT_CAPABILITIES) + retained = len(UNRELEASED_MEMORY) + with pytest.raises(outcome): + _run_check(suite, check) + assert len(UNRELEASED_MEMORY) == retained + 1 + + +def test_a_drain_that_never_returns_fails_the_zero_budget_check_and_is_not_called_again() -> None: + """The check fails at the harness timeout, and the harness then asks nothing more of the + driver: its usual cleanup would drain again and hang the run. The stuck call is released + only once the check is over, so no thread is left behind.""" + release = threading.Event() + budgets: list[float] = [] + + class DrainNeverReturns(MemoryDriver): + def drain(self, budget_s: float) -> bool: + budgets.append(budget_s) + release.wait(TIMEOUT) + return super().drain(budget_s) + + suite = _conformance(DrainNeverReturns, DEFAULT_CAPABILITIES) + failures: list[BaseException] = [] + + def check() -> None: + try: + suite.test_drain_polled_without_a_budget_turns_true_once_nothing_runs() + except BaseException as error: # noqa: BLE001 - judged below + failures.append(error) + + checker = threading.Thread(target=check, daemon=True) + checker.start() + checker.join(TIMEOUT) + finished = not checker.is_alive() + release.set() # lets the stuck drain and, if cleanup drained again, the check return + checker.join(TIMEOUT) + wait_until( + lambda: not any(t.name == OFF_THREAD_NAME for t in threading.enumerate()), + "the released drain call to return", + TIMEOUT, + ) + assert finished, "the check hung in its own cleanup" + assert failures and "did not return" in str(failures[0]), failures + assert budgets == [0.0], "cleanup drained a driver whose drain had not returned" + + class TestMemoryDriverOnChange: @pytest.mark.parametrize("stop", ["cancel", "drain"]) def test_a_cancel_or_drain_shows_before_on_change_is_called(self, stop: str) -> None: