diff --git a/tensorrt_llm/_torch/disaggregation/base/backend.py b/tensorrt_llm/_torch/disaggregation/base/backend.py index a1d0dd5e0e9d..647e66c4db1a 100644 --- a/tensorrt_llm/_torch/disaggregation/base/backend.py +++ b/tensorrt_llm/_torch/disaggregation/base/backend.py @@ -197,8 +197,9 @@ class Cancelled: Outcome = Union[Delivered, Failed, Cancelled] """How one delivery ended. ``None`` rather than a member means it has not ended yet. -Latch which member it is, not the object: ``reports_pending`` turns from true to false over time, so -a stored outcome carries a stale one. +The logical member and its cause are committed at the task/session transition; +polling only observes that decision. ``reports_pending`` can still change, so a +stored outcome carries stale report progress, not a different logical result. ``reports_pending`` asks one question and only one: is a report this transfer was owed still to arrive. Receiving waits on the writers' reports, sending on word about its own writes. @@ -226,6 +227,7 @@ def poll(self) -> Optional[Outcome]: A failure is reported as soon as it is known, which may be before every report is in -- that second question is carried by the outcome itself. + The logical result is stable and does not depend on when the first poll occurs. """ ... diff --git a/tensorrt_llm/_torch/disaggregation/native/handle.py b/tensorrt_llm/_torch/disaggregation/native/handle.py index 0e18e4ddfbfd..1f757d3e1dd5 100644 --- a/tensorrt_llm/_torch/disaggregation/native/handle.py +++ b/tensorrt_llm/_torch/disaggregation/native/handle.py @@ -12,18 +12,10 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""How a native transfer's state reads as a contract outcome. +"""Observe a native task's committed logical result and current report progress. -Both directions land here: a send session and a receive session expose the same three things a -handle needs -- the session's own verdict, its exception, and per-task status. What differs is what -``reports_pending`` is owed by: on the receive side the writers' reports, on the send side word -about this side's own writes. Neither answers whether anyone is still touching the memory; that is a -separate question nothing here asks. - -TODO: Which ending a piece reports depends on when it is first polled. A piece whose session has -gone terminal while the piece itself is still writing latches the session's verdict, yet the piece -may finish afterwards -- so an early poll says failed and a late one says delivered. Nothing polls -these yet; it has to be settled before the surface is frozen. +Task/session transitions commit the result, not polling. Reports and physical +quiescence remain separate questions; an outcome never authorizes memory reuse. """ from __future__ import annotations @@ -32,56 +24,29 @@ from tensorrt_llm._torch.disaggregation.base import Cancelled, Delivered, Failed, Outcome -from .transfer import SessionStatus, TaskStatus +from .transfer import KVRecvTask, KVSendTask, RxSession, SessionStatus, TaskStatus, TxSession class TaskHandle: - """One piece of one request, seen through the session carrying it. - - Pieces of the same request share that session, so this piece's own state is read first and the - session's verdict only answers for a piece that has not ended on its own. - """ + """One piece's decision, including when the first observer arrives late.""" - def __init__(self, session, task, token_end: int): + def __init__( + self, session: RxSession | TxSession, task: KVRecvTask | KVSendTask, token_end: int + ) -> None: self._session = session self._task = task self._token_end = token_end - self._ended: Optional[Outcome] = None def poll(self) -> Optional[Outcome]: - if self._ended is not None: - # Which ending this was cannot change; only whether a writer still owes word can. - return self._rebuild(self._ended) - # This piece's own ending wins: bytes that landed landed, whatever became of its siblings, - # and "cancelled" means stopped short of delivering. - if self._task.status is TaskStatus.TRANSFERRED: - outcome = Delivered(token_end=self._token_end) - elif self._task.status is TaskStatus.ERROR: - if self._ended_by_cancel(): - outcome = Cancelled(by_peer=self._session.cancelled_by_peer, reports_pending=True) - else: - outcome = Failed(reason=self._why_failed(), reports_pending=True) - elif self._session.status is SessionStatus.CANCELLED: - outcome = Cancelled(by_peer=self._session.cancelled_by_peer, reports_pending=True) - elif self._session.status is SessionStatus.ERROR: - outcome = Failed(reason=self._why_failed(), reports_pending=True) - else: + result = self._task.logical_outcome + if result is None: return None - self._ended = outcome - return self._rebuild(outcome) - - def _rebuild(self, ended: Outcome) -> Outcome: - """The latched ending, with today's answer to whether a report is still owed. - - A session is shared by several pieces, so a sibling's failure moves the session's own - verdict after this piece has already ended; the ending latched here does not move with it. - """ - if isinstance(ended, Delivered): - return ended + if result.status is SessionStatus.TRANSFERRED: + return Delivered(token_end=self._token_end) owed = self._owed_a_report() - if isinstance(ended, Cancelled): - return Cancelled(by_peer=ended.by_peer, reports_pending=owed) - return Failed(reason=ended.reason, reports_pending=owed) + if result.status is SessionStatus.CANCELLED: + return Cancelled(by_peer=result.by_peer, reports_pending=owed) + return Failed(reason=result.reason, reports_pending=owed) def _owed_a_report(self) -> bool: """Whether anyone still owes word about this piece. @@ -98,20 +63,6 @@ def _owed_a_report(self) -> bool: return self._task.status is TaskStatus.TRANSFERRING return outstanding - def _ended_by_cancel(self) -> bool: - """Whether this piece's error is the cancellation itself. - - A cancel ends the pieces that had not started by failing them, which is the same task state - a real transfer error leaves; only the recorded cause tells the two apart, and it is the - very object the cancel installed rather than one that merely reads like it. - """ - cancelled = getattr(self._session, "_cancel_exception", None) - return cancelled is not None and self._task._exception is cancelled - - def _why_failed(self) -> str: - error = self._task._exception or self._session.exception - return str(error) if error is not None else "transfer failed without a recorded cause" - class NothingPublished: """Submission failed before any peer was told where to write, so there is nothing to wait on.""" diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index e262fd825ccc..06d4110d1e1c 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -202,6 +202,74 @@ class TaskStatus(Enum): ERROR = "ERROR" +@dataclass(frozen=True) +class _LogicalOutcome: + """A committed delivery decision, independent of physical/report progress.""" + + status: SessionStatus + reason: str = "" + by_peer: bool = False + + +class _LogicalOutcomes: + """Arbitrate task completion against session failure/cancellation at event time. + + A failed or cancelled session ends only pending tasks. Already committed + results, including their original cause, never change. No task, request or + backend handle is retained here: memory ownership remains a separate concern. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._terminal: Optional[_LogicalOutcome] = None + self._results: list[Optional[_LogicalOutcome]] = [] + + def add_task(self) -> int: + with self._lock: + index = len(self._results) + self._results.append(self._terminal) + return index + + def get(self, index: int) -> Optional[_LogicalOutcome]: + with self._lock: + return self._results[index] + + def complete(self, index: int) -> None: + with self._lock: + if self._results[index] is None: + self._results[index] = _LogicalOutcome(SessionStatus.TRANSFERRED) + + def fail(self, error: Exception) -> None: + self._end(_LogicalOutcome(SessionStatus.ERROR, reason=str(error))) + + def cancel(self, by_peer: bool) -> None: + self._end(_LogicalOutcome(SessionStatus.CANCELLED, by_peer=by_peer)) + + def _end(self, outcome: _LogicalOutcome) -> None: + with self._lock: + if self._terminal is not None: + return + self._terminal = outcome + for index, result in enumerate(self._results): + if result is None: + self._results[index] = self._terminal + + +class _LogicalTask: + def __init__(self) -> None: + self._logical_outcomes = _LogicalOutcomes() + self._logical_index = self._logical_outcomes.add_task() + + def bind_logical_outcomes(self, outcomes: _LogicalOutcomes) -> None: + """Join the session before this task is exposed to a worker or caller.""" + self._logical_outcomes = outcomes + self._logical_index = outcomes.add_task() + + @property + def logical_outcome(self) -> Optional[_LogicalOutcome]: + return self._logical_outcomes.get(self._logical_index) + + class _ReceiveOperationOwner: """Track destination access independently from a task's logical result.""" @@ -447,8 +515,9 @@ class _PhysicalOperation: status: Optional[object] = None -class SendTaskBase: +class SendTaskBase(_LogicalTask): def __init__(self, params: DisaggregatedParams): + super().__init__() self.status = TaskStatus.INIT self._event = threading.Event() self._exception: Optional[Exception] = None @@ -460,11 +529,13 @@ def __init__(self, params: DisaggregatedParams): self._physical_operations: dict[int, _PhysicalOperation] = {} def fail(self, exc: Exception) -> None: + self._logical_outcomes.fail(exc) self._exception = exc self.status = TaskStatus.ERROR self._event.set() def complete(self) -> None: + self._logical_outcomes.complete(self._logical_index) self.status = TaskStatus.TRANSFERRED self._event.set() @@ -1848,11 +1919,11 @@ def __init__( self._reported_aux_peer_ranks: set[int] = set() self._has_last_slice = False self.lock = threading.Lock() + self._logical_outcomes = _LogicalOutcomes() self._exception: Optional[Exception] = None self._closed = False self._terminal_status: Optional[SessionStatus] = None - self._cancel_exception: Optional[Exception] = None self.transfer_start_time = None self.transfer_end_time = None # Must be last: makes session visible to listener thread, @@ -1912,6 +1983,7 @@ def send(self, chunk: Chunk) -> None: prompt_len=self._base_args.prompt_len, ) task._unique_rid = self.disagg_request_id + task.bind_logical_outcomes(self._logical_outcomes) self.kv_tasks.append(task) self._has_last_slice |= chunk.is_last req_info_snapshot = dict(self._sender._get_req_info(task._unique_rid) or {}) @@ -1936,6 +2008,7 @@ def send_aux(self) -> AuxSendTask: params = self._base_args.params task = AuxSendTask(params, self.aux_slot) task._unique_rid = self.disagg_request_id + task.bind_logical_outcomes(self._logical_outcomes) self.aux_task = task self._report_unsubmitted_aux_failures(aux_failures) if terminal_error is not None: @@ -2035,19 +2108,16 @@ def cancel(self, by_peer: bool = False) -> bool: def cancel_local(self, by_peer: bool = False) -> bool: aux_failures: list[RecvReqInfo] = [] with self.lock: - # Only an earlier cancel refuses: a failed session may still have peers touching memory, - # and they are told to stop here. Which ending a piece reports is latched by its handle, - # so it does not depend on this. + # A later cancellation must still notify peers after logical failure. + # The logical arbiter preserves whichever outcome already committed. if self._terminal_status == SessionStatus.CANCELLED: return False + self._logical_outcomes.cancel(by_peer) self._terminal_status = SessionStatus.CANCELLED # Who asked is not recoverable later, and the two differ: the peer asking is a transfer # error, our own side asking is an ordinary end. self.cancelled_by_peer = by_peer exc = RuntimeError(f"TxSession {self.disagg_request_id} cancelled") - # Kept so a task failed by the line below is recognisable as cancelled rather than - # broken; its own status cannot say which, since both end it the same way. - self._cancel_exception = exc for task in self.kv_tasks: if task.status == TaskStatus.INIT: task.fail(exc) @@ -2159,6 +2229,7 @@ def wait_for_task(task: SendTaskBase) -> Optional[WaitResult]: self._exception = RuntimeError( "required auxiliary transfer was not dispatched" ) + self._logical_outcomes.fail(self._exception) self._terminal_status = SessionStatus.ERROR return WaitResult.FAILED result = wait_for_task(self.aux_task) @@ -2177,6 +2248,7 @@ def set_exception(self, reason: str = "") -> None: aux_failures: list[RecvReqInfo] = [] with self.lock: self._exception = RuntimeError(msg) + self._logical_outcomes.fail(self._exception) self._terminal_status = SessionStatus.ERROR for task in self.kv_tasks: if not task.is_done: @@ -2219,7 +2291,7 @@ def __del__(self): logger.warning(f"TxSession.__del__: exception during close: {e}") -class KVRecvTask: +class KVRecvTask(_LogicalTask): def __init__( self, unique_rid: Optional[int], @@ -2228,6 +2300,7 @@ def __init__( params: DisaggregatedParams, aux_slot: Optional[int], ): + super().__init__() self._event = threading.Event() self.slice_id = slice_id self.status = TaskStatus.INIT @@ -2245,6 +2318,7 @@ def __init__( self._ownership_state_lock: Optional[threading.Lock] = None def fail(self, exc: Exception) -> None: + self._logical_outcomes.fail(exc) if self._ownership_state_lock is None: self._exception = exc self.status = TaskStatus.ERROR @@ -2257,12 +2331,14 @@ def fail(self, exc: Exception) -> None: def complete(self) -> None: if self._ownership_state_lock is None: + self._logical_outcomes.complete(self._logical_index) self.status = TaskStatus.TRANSFERRED self._event.set() return with self._ownership_state_lock: if self.status == TaskStatus.ERROR: return + self._logical_outcomes.complete(self._logical_index) self.status = TaskStatus.TRANSFERRED self._event.set() @@ -2949,7 +3025,6 @@ def __init__( self._exception: Optional[Exception] = None self._closed = False self._terminal_status: Optional[SessionStatus] = None - self._cancel_exception: Optional[Exception] = None self.transfer_start_time = None self.transfer_end_time = None self.kv_cache_size_bytes: int = 0 @@ -2968,6 +3043,7 @@ def __init__( # publication wins the state transition. self._publication_lock = threading.Lock() self.lock = threading.Lock() + self._logical_outcomes = _LogicalOutcomes() try: self._receiver.setup_session(self) except Exception: @@ -3013,6 +3089,7 @@ def _poison_receiver_ownership(self, error: Exception) -> None: def _record_ownership_evidence_error(self, error: Exception) -> None: """Record a fatal ownership-evidence error and close receiver admission.""" + self._logical_outcomes.fail(error) self._exception = error if self._terminal_status is None: self._terminal_status = SessionStatus.ERROR @@ -3113,6 +3190,7 @@ def receive(self, chunk: Chunk) -> None: params, aux_slot=self.aux_slot, ) + task.bind_logical_outcomes(self._logical_outcomes) self._kv_tasks.append(task) self._receiver.dispatch_task(task) @@ -3129,6 +3207,7 @@ def prepare_receive(self, chunk: Chunk) -> Optional[KVRecvTask]: params, aux_slot=self.aux_slot, ) + task.bind_logical_outcomes(self._logical_outcomes) task.begin_publication() self._kv_tasks.append(task) return task @@ -3147,6 +3226,7 @@ def dispatch_prepared_receive(self, task: KVRecvTask) -> None: def fail_admission(self, error: Exception) -> None: """Fail logical admission without releasing possibly published destinations.""" with self.lock: + self._logical_outcomes.fail(error) self._exception = error if self._terminal_status is None: self._terminal_status = SessionStatus.ERROR @@ -3259,10 +3339,10 @@ def on_done( instance_name=instance_name, instance_rank=instance_rank, ): - # Runs on the scatter worker thread for the bounced path. Touches only this - # task's own status/_event/_perf_timer (no RxSession.lock, no shared session - # state), so it is lock-free. complete() sets status before _event, keeping - # wait_complete's status-first poll correct. + # Runs on the scatter worker thread for the bounced path. Do not acquire + # RxSession.lock here: the non-bounced path invokes this callback inline + # while already holding it. Task outcome/ownership locks are independent. + # complete() sets status before _event for wait_complete's status-first poll. if self._enforce_physical_ownership: task.finish_local_completion() if not success: @@ -3286,8 +3366,8 @@ def on_done( ) task.complete() # Transfer end for perf/time-sync: only meaningful once every slice has - # landed. Plain attribute write (atomic under the GIL); on_done must stay - # lock-free, and consumers only read it after wait_complete succeeds. + # landed. Plain attribute write (atomic under the GIL); consumers only read + # it after wait_complete succeeds. if all(t.status == TaskStatus.TRANSFERRED for t in self._kv_tasks): self.transfer_end_time = tensorrt_llm.bindings.global_steady_clock_now() logger.debug( @@ -3377,12 +3457,14 @@ def process_aux_agent_result(self, peer_rank: int, status: AgentResult): self._exception = RuntimeError( f"Session {self.request_id} received too many aux transfers" ) + self._logical_outcomes.fail(self._exception) if self._terminal_status is None: self._terminal_status = SessionStatus.ERROR logger.error(str(self._exception)) elif status == AgentResult.FAILED: self._aux_status = TaskStatus.ERROR self._exception = RuntimeError(f"Session {self.request_id} aux transfer failed") + self._logical_outcomes.fail(self._exception) if self._terminal_status is None: self._terminal_status = SessionStatus.ERROR else: @@ -3438,19 +3520,16 @@ def resources_drained(self) -> bool: def cancel_local(self, by_peer: bool = False) -> bool: """Commit cancellation under the same lock used for publication.""" with self.lock: - # Only an earlier cancel refuses: a failed session may still have peers touching memory, - # and they are told to stop here. Which ending a piece reports is latched by its handle, - # so it does not depend on this. + # A later cancellation must still notify peers after logical failure. + # The logical arbiter preserves whichever outcome already committed. if self._terminal_status == SessionStatus.CANCELLED: return False + self._logical_outcomes.cancel(by_peer) self._terminal_status = SessionStatus.CANCELLED # Who asked is not recoverable later, and the two differ: the peer asking is a transfer # error, our own side asking is an ordinary end. self.cancelled_by_peer = by_peer exc = RuntimeError(f"RxSession {self.disagg_request_id} cancelled") - # Kept so a task failed by the loop below is recognisable as cancelled rather than - # broken; its own status cannot say which, since both end it the same way. - self._cancel_exception = exc for task in self._kv_tasks: rid_slice = (self.disagg_request_id, task.slice_id) if task.status == TaskStatus.INIT: diff --git a/tests/unittest/disaggregated/test_peer_fetch.py b/tests/unittest/disaggregated/test_peer_fetch.py index 701b0ea02712..a64f136384cb 100644 --- a/tests/unittest/disaggregated/test_peer_fetch.py +++ b/tests/unittest/disaggregated/test_peer_fetch.py @@ -2,9 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 """How the native backend answers the contract's two questions. -The mapping is the whole point of the adapter: which member of ``Outcome`` a native session state -becomes, and whether a writer is still owed word about this piece. Both are exercised against stub -sessions, because the interesting states are the ones a real transfer reaches only under a race. +The adapter observes committed task outcomes and whether a writer still owes word about a piece. +Stub sessions isolate admission, but task transitions use the native logical outcome arbiter. The last section covers what the transceiver holds once admission returns, since an adapter that starts nothing still decides whether a request stays paired with a session the sweep can retire. @@ -29,7 +28,12 @@ ) from tensorrt_llm._torch.disaggregation.base.transfer import get_unique_rid from tensorrt_llm._torch.disaggregation.native.fetch import PeerFetch -from tensorrt_llm._torch.disaggregation.native.transfer import SessionStatus, TaskStatus +from tensorrt_llm._torch.disaggregation.native.transfer import ( + KVRecvTask, + SessionStatus, + TaskStatus, + _LogicalOutcomes, +) from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 from tensorrt_llm.bindings import LlmRequestState @@ -45,25 +49,12 @@ def _owes_a_report(outcome) -> bool: return outcome is None or outcome.reports_pending -class _StubTask: - """A receive task answers for its writers. - - Without that the handle falls back to the send-side probe, and the fan-in assertions would be - testing the wrong branch. Transferring means no writer has reported yet; a test about a settled - piece says so. - """ - - def __init__(self, status, reports_outstanding=True): - self.status = status - self._exception = None - self.reports_outstanding = reports_outstanding - - class _StubSession: """Only what the adapter reads, in states a real session reaches only under a race.""" def __init__(self): self._kv_tasks = [] + self._logical_outcomes = _LogicalOutcomes() self.status = SessionStatus.INIT self.cancelled_by_peer = False self.exception = None @@ -77,18 +68,27 @@ def __init__(self): def receive(self, chunk): if self.raise_on_receive is not None: if self.append_before_raising: - self._kv_tasks.append(_StubTask(TaskStatus.TRANSFERRING)) + self._append_task(chunk) raise self.raise_on_receive if not self.admits: # A closed or already terminal session returns without taking a task. return - self._kv_tasks.append(_StubTask(TaskStatus.TRANSFERRING)) + self._append_task(chunk) + + def _append_task(self, chunk): + task = KVRecvTask(42, chunk, len(self._kv_tasks), DisaggregatedParams(), aux_slot=None) + task.bind_logical_outcomes(self._logical_outcomes) + task.status = TaskStatus.TRANSFERRING + task.expected_transfers = 1 + self._kv_tasks.append(task) def fail_admission(self, error): + self._logical_outcomes.fail(error) self.status = SessionStatus.ERROR self.exception = error def cancel_local(self, by_peer=False): + self._logical_outcomes.cancel(by_peer) self.cancel_committed += 1 return self.cancel_committed == 1 @@ -160,7 +160,7 @@ def test_pieces_of_one_request_share_a_session(): def test_each_attempt_holds_its_own_piece(): worker, peer, first = _fetch_one() second = peer.fetch(_extent()) - worker.session._kv_tasks[0].status = TaskStatus.TRANSFERRED + worker.session._kv_tasks[0].complete() assert isinstance(first.poll(), Delivered) assert second.poll() is None @@ -187,8 +187,8 @@ def test_no_conclusion_while_the_piece_is_in_flight(): def test_delivery_reports_how_far_the_piece_reaches(): worker, _, attempt = _fetch_one(end=64) - worker.session._kv_tasks[0].status = TaskStatus.TRANSFERRED - worker.session._kv_tasks[0].reports_outstanding = False + worker.session._kv_tasks[0].complete() + worker.session._kv_tasks[0].note_writer_report(0, True) outcome = attempt.poll() assert isinstance(outcome, Delivered) assert outcome.token_end == 64 @@ -198,8 +198,7 @@ def test_delivery_reports_how_far_the_piece_reaches(): def test_failure_is_reported_before_every_writer_has_gone_quiet(): """The gate has to stay shut on a failure whose writers have not reported.""" worker, _, attempt = _fetch_one() - worker.session.status = SessionStatus.ERROR - worker.session.exception = RuntimeError("peer died") + worker.session.fail_admission(RuntimeError("peer died")) outcome = attempt.poll() assert isinstance(outcome, Failed) assert "peer died" in outcome.reason @@ -209,27 +208,25 @@ def test_failure_is_reported_before_every_writer_has_gone_quiet(): def test_a_settled_failure_opens_the_gate(): worker, _, attempt = _fetch_one() - worker.session.status = SessionStatus.ERROR - worker.session._kv_tasks[0].status = TaskStatus.ERROR - worker.session._kv_tasks[0].reports_outstanding = False + worker.session.fail_admission(RuntimeError("transfer failed")) + worker.session._kv_tasks[0].note_writer_report(0, False) assert _owes_a_report(attempt.poll()) is False def test_a_cancellation_says_who_asked(): worker, _, attempt = _fetch_one() - worker.session.status = SessionStatus.CANCELLED - worker.session.cancelled_by_peer = True + worker.session.cancel_local(by_peer=True) outcome = attempt.poll() assert isinstance(outcome, Cancelled) assert outcome.by_peer is True assert outcome.reports_pending is True -def test_failure_without_a_recorded_cause_still_says_something(): +def test_failure_preserves_its_recorded_cause(): worker, _, attempt = _fetch_one() - worker.session.status = SessionStatus.ERROR + worker.session.fail_admission(RuntimeError("transfer failed")) assert isinstance(attempt.poll(), Failed) - assert attempt.poll().reason + assert attempt.poll().reason == "transfer failed" # --------------------------------------------------------------------------- diff --git a/tests/unittest/disaggregated/test_task_handle.py b/tests/unittest/disaggregated/test_task_handle.py index 99380759618a..626a5d619cf6 100644 --- a/tests/unittest/disaggregated/test_task_handle.py +++ b/tests/unittest/disaggregated/test_task_handle.py @@ -83,18 +83,23 @@ def _stub_receiver(): return receiver -def _receiving_from(writers: int, rid: int = 7) -> tuple[RxSession, TaskHandle]: +def _receiving_from( + writers: int, rid: int = 7, *, owns_transfers: bool = False +) -> tuple[RxSession, TaskHandle]: """One piece published to ``writers`` peers, and the handle the caller polls it through.""" + receiver = _stub_receiver() + receiver._enforce_physical_ownership = owns_transfers session = RxSession( request_id=rid, params=DisaggregatedParams(disagg_request_id=rid), - receiver=_stub_receiver(), + receiver=receiver, prompt_len=TOKENS, ) session.receive(_sole_piece()) task = session._kv_tasks[0] task.expected_transfers = writers - session.mark_transferring(task.slice_id) + cohort = set(range(writers)) if owns_transfers else None + session.mark_transferring(task.slice_id, cohort) return session, TaskHandle(session, task, TOKENS) @@ -255,20 +260,16 @@ def test_closing_leaves_the_tasks_where_they_were(): assert handle.poll() is None -def test_a_task_that_counts_nothing_falls_back_to_its_own_state(): - """With no count to read, the handle reads the state it does have.""" - session = SimpleNamespace( - status=SessionStatus.ERROR, - exception=RuntimeError("peer died"), - cancelled_by_peer=False, - ) - task = SimpleNamespace(status=TaskStatus.TRANSFERRING, _exception=None) - handle = TaskHandle(session, task, TOKENS) +def test_report_progress_remains_separate_from_a_committed_failure(): + session, handle = _receiving_from(1) + task = session._kv_tasks[0] + task.fail(RuntimeError("peer died")) assert handle.poll().reports_pending is True - task.status = TaskStatus.ERROR + _report(session, 0, AgentResult.SUCCESS) assert handle.poll().reports_pending is False + assert isinstance(handle.poll(), Failed) def test_a_send_task_owes_until_every_peer_write_is_done(): @@ -351,7 +352,7 @@ def test_a_siblings_failure_does_not_reopen_a_delivered_piece(): """Pieces share a session, so its verdict moves after one of them has already ended.""" task = KVRecvTask(9, _sole_piece(), 0, DisaggregatedParams(disagg_request_id=9), aux_slot=None) task.expected_transfers = 1 - task.status = TaskStatus.TRANSFERRED + task.complete() task.note_writer_report(0, True) session = SimpleNamespace( status=SessionStatus.TRANSFERRED, @@ -479,7 +480,7 @@ def test_a_piece_that_landed_is_delivered_even_if_the_request_was_cancelled(): ) task.expected_transfers = 1 task.note_writer_report(0, True) - task.status = TaskStatus.TRANSFERRED + task.complete() session = SimpleNamespace( status=SessionStatus.CANCELLED, exception=None, @@ -636,7 +637,225 @@ def test_a_send_session_reads_a_parked_cancel_as_the_peers_too(): ) assert session.status is SessionStatus.CANCELLED - task = SimpleNamespace(status=TaskStatus.TRANSFERRING, _exception=None) + sender.dispatch_task = MagicMock() + session.send(_sole_piece()) + task = session.kv_tasks[0] outcome = TaskHandle(session, task, TOKENS).poll() assert isinstance(outcome, Cancelled) assert outcome.by_peer is True + + +# --------------------------------------------------------------------------- +# Logical decisions belong to the transition, not the first observer +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("poll_before_completion", [False, True]) +@pytest.mark.parametrize("ending", ["failure", "local_cancel", "peer_cancel"]) +def test_receive_outcome_does_not_depend_on_polling_before_late_completion( + ending: str, poll_before_completion: bool +) -> None: + session, handle = _receiving_from(1) + if ending == "failure": + session.fail_admission(RuntimeError("another publication failed")) + else: + session.cancel_local(by_peer=ending == "peer_cancel") + if poll_before_completion: + assert handle.poll() is not None + + # The writer finishes after the logical decision. Its report still drains the transfer. + _report(session, 0, AgentResult.SUCCESS) + + for observer in (handle, TaskHandle(session, session._kv_tasks[0], TOKENS)): + outcome = observer.poll() + if ending == "failure": + assert isinstance(outcome, Failed) + assert "another publication failed" in outcome.reason + else: + assert isinstance(outcome, Cancelled) + assert outcome.by_peer is (ending == "peer_cancel") + assert outcome.reports_pending is False + + +@pytest.mark.parametrize("by_peer", [False, True]) +def test_receive_cancel_keeps_its_outcome_when_a_writer_later_fails(by_peer: bool) -> None: + session, _ = _receiving_from(1) + session.cancel_local(by_peer=by_peer) + + _report(session, 0, AgentResult.FAILED) + + outcome = TaskHandle(session, session._kv_tasks[0], TOKENS).poll() + assert isinstance(outcome, Cancelled) + assert outcome.by_peer is by_peer + assert outcome.reports_pending is False + + +def test_receive_session_failure_precedes_later_cancel_without_an_observer() -> None: + session, handle = _receiving_from(1) + session.fail_admission(RuntimeError("first publication failure")) + + assert session.cancel_local() is True + + outcome = handle.poll() + assert isinstance(outcome, Failed) + assert "first publication failure" in outcome.reason + assert outcome.reports_pending is True + + +def _sending_pieces(count: int = 1) -> tuple[Sender, TxSession]: + sender = _wired_sender() + sender.dispatch_task = MagicMock() + sender._get_result_dealer = MagicMock() + sender._instance_rank = 0 + session = TxSession( + request_id=30, params=DisaggregatedParams(disagg_request_id=30), sender=sender + ) + for _ in range(count): + session.send(_sole_piece()) + return sender, session + + +@pytest.mark.parametrize("by_peer", [False, True]) +@pytest.mark.parametrize("queued_status", [TaskStatus.INIT, TaskStatus.TRANSFERRING]) +def test_queued_sender_abort_preserves_the_committed_cancellation( + by_peer: bool, queued_status: TaskStatus +) -> None: + sender, session = _sending_pieces() + task = session.kv_tasks[0] + task.status = queued_status + session.cancel_local(by_peer=by_peer) + # Execute the real worker's pre-submission abort branch after cancellation won. + empty = SimpleNamespace(size=0) + write_meta = SimpleNamespace( + src_ptrs=empty, + dst_ptrs=empty, + sizes=empty, + unique_rid=30, + slice_id=0, + receiver_slice_id=0, + peer_rank=0, + peer_endpoint="tcp://receiver:1234", + task=task, + ) + + sender._deliver_kv_to_agent(write_meta) + + outcome = TaskHandle(session, task, TOKENS).poll() + assert isinstance(outcome, Cancelled) + assert outcome.by_peer is by_peer + sender._get_result_dealer.return_value.send.assert_called_once() + + +@pytest.mark.parametrize("poll_before_completion", [False, True]) +def test_sender_sibling_failure_is_stable_across_late_completion( + poll_before_completion: bool, +) -> None: + _, session = _sending_pieces(2) + failed, pending = session.kv_tasks + pending.status = TaskStatus.TRANSFERRING + handle = TaskHandle(session, pending, TOKENS) + + failed.fail(RuntimeError("first sibling failed")) + if poll_before_completion: + assert isinstance(handle.poll(), Failed) + pending.complete() + + for observer in (handle, TaskHandle(session, pending, TOKENS)): + outcome = observer.poll() + assert isinstance(outcome, Failed) + assert "first sibling failed" in outcome.reason + + +def test_terminal_failure_cause_is_stable_without_polling() -> None: + session, handle = _receiving_from(1) + task = session._kv_tasks[0] + task.fail(RuntimeError("original failure")) + task.fail(RuntimeError("cleanup failure")) + + outcome = handle.poll() + assert isinstance(outcome, Failed) + assert outcome.reason == "original failure" + + +@pytest.mark.parametrize("ending", ["failure", "local_cancel", "peer_cancel"]) +def test_delivered_outcome_precedes_later_session_terminal_events(ending: str) -> None: + session, handle = _receiving_from(1) + _report(session, 0, AgentResult.SUCCESS) + if ending == "failure": + session.fail_admission(RuntimeError("later failure")) + else: + session.cancel_local(by_peer=ending == "peer_cancel") + + assert isinstance(handle.poll(), Delivered) + assert handle.poll().token_end == TOKENS + + +def test_delivered_sender_piece_survives_a_direct_sibling_failure_without_polling() -> None: + _, session = _sending_pieces(2) + delivered, failed = session.kv_tasks + delivered.complete() + + failed.fail(RuntimeError("sibling failed later")) + + assert isinstance(TaskHandle(session, delivered, TOKENS).poll(), Delivered) + assert isinstance(TaskHandle(session, failed, TOKENS).poll(), Failed) + + +@pytest.mark.parametrize("owns_transfers", [False, True]) +@pytest.mark.parametrize("scatter_succeeded", [False, True]) +@pytest.mark.parametrize("cancel_first", [False, True]) +def test_scatter_callback_and_cancel_commit_in_event_order( + monkeypatch: pytest.MonkeyPatch, + owns_transfers: bool, + scatter_succeeded: bool, + cancel_first: bool, +) -> None: + from tensorrt_llm._torch.disaggregation.native import bounce + + session, handle = _receiving_from(1, owns_transfers=owns_transfers) + deferred = [] + monkeypatch.setattr(bounce, "scatter_write_result", lambda *args: deferred.append(args[-1])) + _report(session, 0, AgentResult.SUCCESS) + assert len(deferred) == 1 + assert handle.poll() is None + if owns_transfers: + assert session.resources_drained() is False + + ready, resume, finished = threading.Event(), threading.Event(), threading.Event() + + def finish_scatter() -> None: + ready.set() + if resume.wait(timeout=5): + deferred[0](scatter_succeeded) + finished.set() + + worker = threading.Thread(target=finish_scatter) + worker.start() + try: + assert ready.wait(timeout=5) + if cancel_first: + session.cancel_local(by_peer=True) + assert isinstance(handle.poll(), Cancelled) + if owns_transfers: + assert session.resources_drained() is False + resume.set() + assert finished.wait(timeout=5) + if not cancel_first: + session.cancel_local(by_peer=True) + finally: + resume.set() + worker.join(timeout=5) + assert not worker.is_alive() + + for observer in (handle, TaskHandle(session, session._kv_tasks[0], TOKENS)): + outcome = observer.poll() + if cancel_first: + assert isinstance(outcome, Cancelled) + assert outcome.by_peer is True + elif scatter_succeeded: + assert isinstance(outcome, Delivered) + else: + assert isinstance(outcome, Failed) + assert "bounce scatter failed" in outcome.reason + if owns_transfers: + assert session.resources_drained() is True diff --git a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py index 0e4bb46daa84..e61a416baa68 100644 --- a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py +++ b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py @@ -36,6 +36,7 @@ TransferWorker, TransferWorkerConfig, TxSession, + _LogicalOutcomes, ) from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 from tensorrt_llm.bindings import LlmRequestState @@ -171,6 +172,7 @@ def _make_tx_session( deadline_monotonic_s: Optional[float] = None, ) -> TxSession: session = object.__new__(TxSession) + session._logical_outcomes = _LogicalOutcomes() session._timeout_s = timeout_s session._overall_timeout_s = None session._deadline_monotonic_s = deadline_monotonic_s diff --git a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py index 583204fc260a..5c6119cd6a07 100644 --- a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py +++ b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py @@ -1601,12 +1601,14 @@ def test_terminal_sender_settles_known_unsubmitted_aux_peer_once() -> None: session._enforce_physical_ownership = True session._need_aux = True session._reported_aux_peer_ranks = set() + session._logical_outcomes = transfer_mod._LogicalOutcomes() session.kv_tasks = [] session.lock = threading.Lock() session._exception = None session._terminal_status = None session._closed = False session.aux_task = transfer_mod.AuxSendTask(params, slot=0) + session.aux_task.bind_logical_outcomes(session._logical_outcomes) assert session.aux_task.begin_physical_operation(active_info.instance_rank) session.set_exception("request failed before auxiliary submission")