[None][fix] Guard Python KV receive ownership and publication

tensorrt_llm/_torch/disaggregation/native/bounce/core.py

@@ -75,6 +75,8 @@ class TransferContext:
     _writer_ok: Dict[int, bool] = field(default_factory=dict)
     # per successful writer: where it wrote, plus the fragments to scatter back
     _scatter_descs: List[tuple] = field(default_factory=list)
+    _writer_cohort: Optional[frozenset[int]] = None
+    _publication_failed: bool = False
     _orphaned: bool = False
     scatter_state: ScatterState = ScatterState.IDLE
     state: TransferState = TransferState.INIT
@@ -93,10 +95,16 @@ def _writers_final(self) -> bool:
         )
 
     def _all_writers_reported(self) -> bool:
+        if self._writer_cohort is not None:
+            return self._writer_cohort.issubset(self._writer_ok)
         return len(self._writer_ok) >= self.num_writers
 
     def _all_writers_succeeded(self) -> bool:
-        return self._all_writers_reported() and all(self._writer_ok.values())
+        if not self._all_writers_reported():
+            return False
+        if self._writer_cohort is not None:
+            return all(self._writer_ok[rank] for rank in self._writer_cohort)
+        return all(self._writer_ok.values())
 
     # Mutations: call only while holding the transport's reservation lock.
     def record_writer_result(
@@ -115,6 +123,8 @@ def record_writer_result(
         # notification, or a stray failure that would flip a good transfer to failed; drop it.
         if self._writers_final() or peer_rank in self._writer_ok:
             return
+        if self._writer_cohort is not None and peer_rank not in self._writer_cohort:
+            raise RuntimeError(f"writer {peer_rank} is outside the published bounce cohort")
         self._writer_ok[peer_rank] = succeeded
         if succeeded and dst_ptrs is not None and int(dst_ptrs.size) > 0:
             self._scatter_descs.append(
@@ -130,6 +140,25 @@ def mark_orphaned(self) -> None:
             return
         self._orphaned = True
 
+    def abort_publication(self, published_writers: set[int]) -> None:
+        """Close a failed fan-out around only the writers that were published.
+
+        Published writers still have to report terminal evidence. Successful
+        partial data is intentionally not scattered because the request is
+        already incomplete.
+        """
+        if self._writers_final():
+            return
+        published = frozenset(published_writers)
+        if len(published) > self.num_writers or not self._writer_ok.keys() <= published:
+            raise RuntimeError(
+                "published bounce cohort is inconsistent with terminal evidence: "
+                f"reported={sorted(self._writer_ok)} published={sorted(published)} "
+                f"reserved_count={self.num_writers}"
+            )
+        self._writer_cohort = published
+        self._publication_failed = True
+
     def begin_scatter(self) -> None:
         self.state = TransferState.SCATTERING
         self.scatter_state = ScatterState.QUEUED
@@ -146,6 +175,7 @@ def ready_to_scatter(self) -> bool:
         return (
             self.state is TransferState.ACTIVE
             and not self._orphaned
+            and not self._publication_failed
             and self._all_writers_succeeded()
             and bool(self._scatter_descs)
         )
@@ -157,6 +187,8 @@ def ready_to_settle(self) -> bool:
             return True  # in doubt: settle now and quarantine
         if not self._all_writers_reported():
             return False  # a writer has not reported yet
+        if self._publication_failed:
+            return True  # every published writer drained; incomplete data is discarded
         if self._all_writers_succeeded() and self._scatter_descs:
             return self.scatter_state in (ScatterState.DONE, ScatterState.FAILED)
         return True  # nothing to scatter, or a failure among them: either way drained
@@ -170,7 +202,11 @@ def settle(self) -> Optional[Settlement]:
         if self._orphaned:
             self.state = TransferState.QUARANTINED
             return Settlement(self.slot_id, Disposition.QUARANTINE, False, self.on_done)
-        success = self._all_writers_succeeded() and self.scatter_state is not ScatterState.FAILED
+        success = (
+            not self._publication_failed
+            and self._all_writers_succeeded()
+            and self.scatter_state is not ScatterState.FAILED
+        )
         self.state = TransferState.COMPLETED if success else TransferState.FAILED
         return Settlement(self.slot_id, Disposition.RELEASE, success, self.on_done)
 
@@ -220,6 +256,10 @@ def release_idle_reservation(self, rid_slice) -> None:
     def orphan_reservation(self, rid_slice) -> None:
         """Give up on an in-flight reservation (cancel/timeout/lost result); quarantine, don't leak."""
 
+    @abstractmethod
+    def abort_publication(self, rid_slice, published_writers: set[int]) -> None:
+        """Limit a failed fan-out to writers whose REQUEST_DATA was queued."""
+
     @abstractmethod
     def record_result(
         self, rid_slice, peer_rank, dst_ptrs=None, sizes=None, src_base=None, on_done=None

tensorrt_llm/_torch/disaggregation/native/bounce/impl.py

@@ -379,6 +379,13 @@ def orphan_reservation(self, rid_slice: RidSlice) -> None:
         or leaking it. Idempotent; a no-op once the transfer has settled."""
         self._apply(rid_slice, lambda ctx: ctx.mark_orphaned())
 
+    def abort_publication(self, rid_slice: RidSlice, published_writers: set[int]) -> None:
+        """Retain a failed fan-out until every successfully published writer drains."""
+        self._apply(
+            rid_slice,
+            lambda ctx: ctx.abort_publication(published_writers),
+        )
+
     def _apply(self, rid_slice: RidSlice, mutate: Callable[[TransferContext], None]) -> None:
         """Mutate the state under the lock, then do what it asks (scatter or settle) with the lock
         released, never holding it across a CUDA sync, a queue put, or a callback. No-op if the
@@ -527,6 +534,9 @@ def release_idle_reservation(self, rid_slice) -> None:
     def orphan_reservation(self, rid_slice) -> None:
         pass
 
+    def abort_publication(self, rid_slice, published_writers: set[int]) -> None:
+        pass
+
     def record_result(
         self, rid_slice, peer_rank, dst_ptrs=None, sizes=None, src_base=None, on_done=None
     ):