[None][test] Improve store driver conformance for fixed pools and shutdown

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

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

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",
+    ),
 }