From e4a3411c98f6d183c679d38004b0d5175122f881 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 2 Oct 2026 14:25:07 -0700 Subject: [PATCH] Extract shared KV transfer lifecycle enforcement primitives Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../_torch/disaggregation/lifecycle/README.md | 36 ++ .../disaggregation/lifecycle/__init__.py | 16 + .../disaggregation/lifecycle/ownership.py | 490 ++++++++++++++++++ .../disaggregation/lifecycle/retirement.py | 346 +++++++++++++ .../disaggregation/native/retirement.py | 362 +------------ .../_torch/disaggregation/native/transfer.py | 458 ++-------------- .../disaggregated/test_lifecycle_binding.py | 248 +++++++++ 7 files changed, 1195 insertions(+), 761 deletions(-) create mode 100644 tensorrt_llm/_torch/disaggregation/lifecycle/README.md create mode 100644 tensorrt_llm/_torch/disaggregation/lifecycle/__init__.py create mode 100644 tensorrt_llm/_torch/disaggregation/lifecycle/ownership.py create mode 100644 tensorrt_llm/_torch/disaggregation/lifecycle/retirement.py create mode 100644 tests/unittest/disaggregated/test_lifecycle_binding.py diff --git a/tensorrt_llm/_torch/disaggregation/lifecycle/README.md b/tensorrt_llm/_torch/disaggregation/lifecycle/README.md new file mode 100644 index 000000000000..de60058b03da --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/lifecycle/README.md @@ -0,0 +1,36 @@ +# Shared lifecycle enforcement + +This internal package reuses the native Python transfer lifecycle implementation. +It does not introduce another KV owner, allocator, public backend API, or supported +configuration. Native sessions still select writers, order publication, submit +transfers, translate backend evidence, and synchronize local CUDA work. + +## Binding responsibilities + +| Mechanism | Shared implementation | Integration responsibility | +| --- | --- | --- | +| Source access | `SendOperationOwner` retains the exact request/status across ambiguous completion. | Map participants and root allocation/registration loans before submission; supply backend-defined completion. | +| Destination access | `ReceiveOperationOwner` accounts for the sealed writer cohort and local completion. | Select the cohort, serialize publication against cancellation, and retain destination storage. A receive owner tracks evidence, not the allocation itself. | +| Deadline arbitration | `RetirementDeadline` serializes exposure, evidence, request timeout, and grace expiry. | Bind all pieces of one session to the same controller; seal completion only after all expected pieces, including AUX, are accounted for. | +| Containment | `RetirementWatchdog` closes admission and notifies containment once. | Keep sessions, manager loans, pools, and agents reachable; provide the qualified executor containment callback. | +| Logical outcome | Remains with the integration, outside this package. | Commit failure/cancellation once; bind a metadata-only timeout callback. Late physical completion must not replace that outcome. | + +Only `NOT_SUBMITTED` and `BACKEND_DONE` establish source access-end. The other +states remain `ADMITTED`, `SUBMITTING`, `SUBMITTED`, and `IN_DOUBT`. A backend +error, elapsed time, logical failure, or empty report count is not quiescence. +Retirement at or after fatal expiry remains prohibited even if DONE arrives later. + +The integration must couple evidence and strong resource roots. Merely constructing +an operation owner does not pin KV Manager pages, hold a registration, or acquire a +manager reference. `expose(*owners)` can retain multiple opaque claims atomically; +each claim requires its own access-end evidence before `settle()`. In particular, +`resources_drained` includes the shared session predicate, so a separately exposed +loan cannot wait on that predicate to settle itself. Release/refcount operations +remain with the resource owner and must be idempotent. + +Native compatibility imports refer to these same classes; native tasks delegate +source bookkeeping without changing wire messages, activation, or task outcomes. +The module-level binding tests exercise opaque resource holds and completion +probes, not a production runtime adapter or a qualified shared-transfer profile. +The runtime integration must still agree its participant mapping, session boundary, +resource-root binding, backend evidence, and containment callback before adoption. diff --git a/tensorrt_llm/_torch/disaggregation/lifecycle/__init__.py b/tensorrt_llm/_torch/disaggregation/lifecycle/__init__.py new file mode 100644 index 000000000000..17f3c426462d --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/lifecycle/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Internal lifecycle enforcement; no public resource or transfer API.""" diff --git a/tensorrt_llm/_torch/disaggregation/lifecycle/ownership.py b/tensorrt_llm/_torch/disaggregation/lifecycle/ownership.py new file mode 100644 index 000000000000..2c9327415b83 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/lifecycle/ownership.py @@ -0,0 +1,490 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Internal physical-access bookkeeping shared by transfer integrations. + +This is not an allocator or a public cache-backend API. Callers supply already +mapped participants and backend-defined access-end evidence, and retain manager +loans until the deadline permits retirement. Logical Attempt outcomes remain +separate. Native publication ordering, transport submission and local-copy +synchronization are not implemented here. +""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING, Optional + +from .retirement import RetirementDeadline + +if TYPE_CHECKING: + from ..base.agent import TransferStatus + + +class TransferNotSubmittedError(RuntimeError): + """Admission rejected an operation before the backend could access memory.""" + + +class PhysicalOperationState(Enum): + """Sender-side evidence for one peer's access to a task's source memory. + + The legal forward paths are:: + + ADMITTED -> NOT_SUBMITTED + ADMITTED -> SUBMITTING -> SUBMITTED -> BACKEND_DONE + SUBMITTING -> IN_DOUBT + SUBMITTED -> IN_DOUBT -> BACKEND_DONE (same retained status reports DONE) + + Only NOT_SUBMITTED and BACKEND_DONE prove that the operation can no longer + access the source. An IN_DOUBT operation without a retained backend status + cannot retire. Repeating a terminal transition is idempotent, and repeating + IN_DOUBT preserves the retained backend evidence. + """ + + ADMITTED = "ADMITTED" + SUBMITTING = "SUBMITTING" + SUBMITTED = "SUBMITTED" + NOT_SUBMITTED = "NOT_SUBMITTED" + BACKEND_DONE = "BACKEND_DONE" + IN_DOUBT = "IN_DOUBT" + + +_DRAINED_PHYSICAL_OPERATION_STATES = frozenset( + (PhysicalOperationState.NOT_SUBMITTED, PhysicalOperationState.BACKEND_DONE) +) + + +@dataclass +class _PhysicalOperation: + """State plus strong backend roots retained until physical quiescence.""" + + state: PhysicalOperationState + request: Optional[object] = None + status: Optional[TransferStatus] = None + + +class SendOperationOwner: + """Retain one source operation per participant independently of logical results. + + The caller maps participant IDs and retains allocation/registration loans. + This helper owns only the operation request and completion-status roots. + Backend submission, evidence translation and resource release stay with + the caller; an ambiguous error is never safe completion. + """ + + def __init__(self, retirement: Optional[RetirementDeadline] = None) -> None: + """Create bookkeeping before admitting any physical operations. + + Args: + retirement: Optional session deadline sharing the admission arbiter. + """ + self._retirement = retirement + self._physical_lock = threading.Lock() if retirement is None else retirement.lock + self._physical_operations: dict[int, _PhysicalOperation] = {} + + def bind_retirement(self, retirement: Optional[RetirementDeadline]) -> None: + """Bind the session arbiter before the owner is exposed to workers. + + Args: + retirement: Deadline used by every operation in the session. + """ + self._retirement = retirement + if retirement is not None: + self._physical_lock = retirement.lock + + def begin_physical_operation(self, peer_rank: int) -> bool: + with self._physical_lock: + if peer_rank in self._physical_operations: + return False + self._physical_operations[peer_rank] = _PhysicalOperation( + PhysicalOperationState.ADMITTED + ) + return True + + def _require_physical_operation_locked( + self, + peer_rank: int, + expected_states: tuple[PhysicalOperationState, ...], + ) -> _PhysicalOperation: + operation = self._physical_operations.get(peer_rank) + if operation is None: + raise RuntimeError(f"physical operation {peer_rank} was not admitted") + if operation.state not in expected_states: + expected = ", ".join(state.value for state in expected_states) + raise RuntimeError( + f"physical operation {peer_rank} is {operation.state.value}, expected {expected}" + ) + return operation + + def begin_backend_submission( + self, + peer_rank: int, + request: object, + ) -> None: + with self._physical_lock: + operation = self._require_physical_operation_locked( + peer_rank, (PhysicalOperationState.ADMITTED,) + ) + if self._retirement is not None and not self._retirement.expose(operation): + operation.state = PhysicalOperationState.NOT_SUBMITTED + raise TransferNotSubmittedError("source retirement admission is closed") + operation.request = request + operation.state = PhysicalOperationState.SUBMITTING + + def record_backend_submission(self, peer_rank: int, status: TransferStatus) -> None: + with self._physical_lock: + operation = self._require_physical_operation_locked( + peer_rank, (PhysicalOperationState.SUBMITTING,) + ) + operation.status = status + operation.state = PhysicalOperationState.SUBMITTED + + def mark_physical_operation_in_doubt(self, peer_rank: int) -> None: + with self._physical_lock: + operation = self._require_physical_operation_locked( + peer_rank, + ( + PhysicalOperationState.SUBMITTING, + PhysicalOperationState.SUBMITTED, + PhysicalOperationState.IN_DOUBT, + ), + ) + operation.state = PhysicalOperationState.IN_DOUBT + if self._retirement is not None: + self._retirement.request_drain("backend quiescence unproven") + + def retire_unsubmitted_physical_operation(self, peer_rank: int) -> None: + with self._physical_lock: + operation = self._require_physical_operation_locked( + peer_rank, + ( + PhysicalOperationState.ADMITTED, + PhysicalOperationState.NOT_SUBMITTED, + ), + ) + operation.state = PhysicalOperationState.NOT_SUBMITTED + + def retire_backend_done_physical_operation(self, peer_rank: int) -> bool: + with self._physical_lock: + operation = self._require_physical_operation_locked( + peer_rank, + ( + PhysicalOperationState.SUBMITTED, + PhysicalOperationState.BACKEND_DONE, + ), + ) + if operation.state is PhysicalOperationState.BACKEND_DONE: + return True + if self._retirement is not None and not self._retirement.settle(operation): + operation.state = PhysicalOperationState.IN_DOUBT + return False + operation.request = None + operation.status = None + operation.state = PhysicalOperationState.BACKEND_DONE + return True + + def poll_in_doubt_physical_operation(self, peer_rank: int) -> bool: + """Retire once, only after a fresh DONE query on the retained status. + + Polling does not change the task's logical outcome. Keep strong local + roots across the query, and reject a result if its operation changed. + """ + with self._physical_lock: + operation = self._physical_operations.get(peer_rank) + if operation is None or operation.state is not PhysicalOperationState.IN_DOUBT: + return False + request, status = operation.request, operation.status + if status is None: + return False + try: + completed = status.is_completed() + except Exception: + # A backend query failure is not evidence that its accessors stopped. + return False + if completed is not True: + return False + with self._physical_lock: + if ( + self._physical_operations.get(peer_rank) is not operation + or operation.state is not PhysicalOperationState.IN_DOUBT + or operation.request is not request + or operation.status is not status + ): + return False + if self._retirement is not None and not self._retirement.settle(operation): + return False + operation.request = None + operation.status = None + operation.state = PhysicalOperationState.BACKEND_DONE + return True + + def has_started_physical_operation(self, peer_rank: int) -> bool: + with self._physical_lock: + return peer_rank in self._physical_operations + + @property + def resources_drained(self) -> bool: + with self._physical_lock: + return (self._retirement is None or self._retirement.can_retire()) and all( + operation.state in _DRAINED_PHYSICAL_OPERATION_STATES + for operation in self._physical_operations.values() + ) + + +class ReceiveOperationOwner: + """Track destination access independently from a task's logical result.""" + + def __init__(self, retirement: Optional[RetirementDeadline] = None) -> None: + self._retirement = retirement + self._lock = threading.Lock() if retirement is None else retirement.lock + self._publication_pending = False + self._cancelled_unpublished = False + self._expected_writers: Optional[int] = None + self._writer_cohort: Optional[frozenset[int]] = None + self._writer_candidates: Optional[frozenset[int]] = None + self._published_writers: Optional[frozenset[int]] = None + self._quiesced_sessions: set[int] = set() + self._writer_results: dict[int, bool] = {} + self._in_doubt_writers: set[int] = set() + self._settled_writers: set[int] = set() + self._publication_failed = False + self._local_completion_pending = False + self._invalid_evidence = False + + def _settle_if_drained_locked(self) -> None: + """Remove this claim only after the complete cohort and local work settle.""" + if self._retirement is not None and self._resources_drained_locked(): + self._retirement.settle(self) + + def _invalidate_evidence_locked(self) -> None: + """Quarantine conflicting proof, including proof disputed after settlement.""" + self._invalid_evidence = True + if self._retirement is not None: + self._retirement.retain_unproven(self, "invalid receive ownership evidence") + + def begin_publication(self) -> None: + with self._lock: + if self._publication_pending or self._expected_writers is not None: + raise RuntimeError("destination publication was already started") + self._publication_pending = True + + def seal_writer_cohort( + self, + expected_writers: int, + writer_cohort: Optional[set[int]] = None, + *, + published_writers: Optional[set[int]] = None, + ) -> None: + if expected_writers < 0: + raise ValueError(f"expected_writers must be non-negative, got {expected_writers}") + cohort = None if writer_cohort is None else frozenset(writer_cohort) + candidates = cohort if published_writers is None else frozenset(published_writers) + if cohort is not None and len(cohort) != expected_writers: + raise ValueError( + f"writer cohort has {len(cohort)} member(s), expected {expected_writers}" + ) + if candidates is not None and ( + len(candidates) < expected_writers or (cohort is not None and not cohort <= candidates) + ): + raise ValueError("published candidates must cover every eligible writer") + with self._lock: + if self._expected_writers is not None: + if ( + self._expected_writers != expected_writers + or self._writer_cohort != cohort + or self._writer_candidates != candidates + ): + raise RuntimeError("writer cohort was already sealed differently") + return + self._expected_writers = expected_writers + self._writer_cohort = cohort + self._writer_candidates = candidates + self._published_writers = candidates + + def finish_publication(self) -> None: + """Record that every authorized REQUEST_DATA message was sent.""" + with self._lock: + self._publication_pending = False + self._settle_if_drained_locked() + + def abort_publication(self, published_writers: set[int]) -> None: + """Close a failed fan-out around the writers whose sends succeeded.""" + with self._lock: + published = frozenset(published_writers) + if self._writer_cohort is not None and not published.issubset(self._writer_cohort): + self._invalidate_evidence_locked() + raise RuntimeError("publication recorded a writer outside the sealed cohort") + if not (self._writer_results.keys() | self._in_doubt_writers) <= published: + self._invalidate_evidence_locked() + raise RuntimeError("terminal evidence came from an unpublished writer") + self._expected_writers = len(published) + self._writer_cohort = frozenset(published) + self._published_writers = published + self._publication_failed = True + self._cancelled_unpublished = not published + self._publication_pending = False + self._settle_if_drained_locked() + + def cancel_unpublished(self) -> bool: + """Close a publication that did not authorize a remote writer.""" + with self._lock: + if self._invalid_evidence: + return False + if self._expected_writers is not None: + return self._cancelled_unpublished + self._cancelled_unpublished = True + self._expected_writers = 0 + self._writer_cohort = frozenset() + self._published_writers = frozenset() + self._publication_pending = False + self._settle_if_drained_locked() + return True + + def record_session_quiesced(self, peer_rank: int) -> None: + """Record no-future-access proof without inventing a per-piece result.""" + with self._lock: + if self._writer_candidates is None or peer_rank not in self._writer_candidates: + self._invalidate_evidence_locked() + raise RuntimeError(f"session acknowledgment from unknown writer {peer_rank}") + if self._published_writers is not None and peer_rank in self._published_writers: + self._quiesced_sessions.add(peer_rank) + self._settle_if_drained_locked() + + def record_writer_in_doubt(self, peer_rank: int) -> bool: + """Retain ownership after a writer reports no safe terminal evidence.""" + with self._lock: + if self._expected_writers is None: + self._invalidate_evidence_locked() + raise RuntimeError( + f"writer {peer_rank} reported ambiguous evidence before publication" + ) + if self._writer_cohort is not None and peer_rank not in self._writer_cohort: + self._invalidate_evidence_locked() + raise RuntimeError(f"writer {peer_rank} is outside the sealed cohort") + if peer_rank in self._in_doubt_writers or peer_rank in self._settled_writers: + return False + self._in_doubt_writers.add(peer_rank) + if self._retirement is not None: + self._retirement.retain_unproven(self, "receive ownership is in doubt") + return True + + def record_writer_settlement(self, peer_rank: int) -> bool: + """Accept physical DONE for a writer that previously reported IN_DOUBT. + + Ordinary FAILED is not this proof: it may describe a later, unsubmitted + chunk while the earlier ambiguous write is still touching the destination. + """ + with self._lock: + if self._expected_writers is None or ( + self._writer_cohort is not None and peer_rank not in self._writer_cohort + ): + self._invalidate_evidence_locked() + raise RuntimeError(f"writer {peer_rank} settled outside the published cohort") + if peer_rank in self._settled_writers: + return False + if peer_rank not in self._in_doubt_writers: + self._invalidate_evidence_locked() + raise RuntimeError(f"writer {peer_rank} settled without prior ambiguous evidence") + if self._writer_results.get(peer_rank) is True: + self._invalidate_evidence_locked() + raise RuntimeError(f"writer {peer_rank} settled after contradictory success") + self._writer_results[peer_rank] = False + self._in_doubt_writers.remove(peer_rank) + self._settled_writers.add(peer_rank) + self._settle_if_drained_locked() + return True + + def record_writer_result( + self, + peer_rank: int, + succeeded: bool, + *, + wait_for_local_completion: bool, + ) -> tuple[bool, bool]: + """Record one writer and return ``(accepted, all_succeeded)``.""" + with self._lock: + if self._expected_writers is None: + self._invalidate_evidence_locked() + raise RuntimeError( + f"writer {peer_rank} reported terminal evidence before publication" + ) + if self._writer_cohort is not None and peer_rank not in self._writer_cohort: + self._invalidate_evidence_locked() + raise RuntimeError(f"writer {peer_rank} is outside the sealed cohort") + if peer_rank in self._in_doubt_writers: + if succeeded: + self._invalidate_evidence_locked() + raise RuntimeError(f"writer {peer_rank} reported success while in doubt") + return False, False + previous = self._writer_results.get(peer_rank) + if previous is not None: + if previous != succeeded: + self._invalidate_evidence_locked() + raise RuntimeError( + f"writer {peer_rank} reported contradictory terminal evidence" + ) + return False, False + if len(self._writer_results) >= self._expected_writers: + return False, False + self._writer_results[peer_rank] = succeeded + all_reported = len(self._writer_results) == self._expected_writers + all_succeeded = ( + all_reported and not self._publication_failed and all(self._writer_results.values()) + ) + if all_succeeded and wait_for_local_completion: + self._local_completion_pending = True + self._settle_if_drained_locked() + return True, all_succeeded + + def finish_local_completion(self) -> None: + with self._lock: + self._local_completion_pending = False + self._settle_if_drained_locked() + + @property + def all_writers_reported(self) -> bool: + """Whether every writer of the sealed cohort has reported a terminal result.""" + with self._lock: + return ( + self._expected_writers is not None + and len(self._writer_results) == self._expected_writers + and not self._in_doubt_writers + ) + + @property + def resources_drained(self) -> bool: + with self._lock: + return self._resources_drained_locked() and ( + self._retirement is None or self._retirement.can_retire() + ) + + def _resources_drained_locked(self) -> bool: + """Check physical evidence without consulting the session's other owners.""" + return ( + self._expected_writers is not None + and ( + len(self._writer_results) == self._expected_writers + or ( + self._published_writers is not None + and self._published_writers <= self._quiesced_sessions + ) + ) + and not self._publication_pending + and not self._local_completion_pending + and not self._in_doubt_writers + and not self._invalid_evidence + ) diff --git a/tensorrt_llm/_torch/disaggregation/lifecycle/retirement.py b/tensorrt_llm/_torch/disaggregation/lifecycle/retirement.py new file mode 100644 index 000000000000..455376d42905 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/lifecycle/retirement.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Deadline arbitration that never queries or releases a transfer backend.""" + +from __future__ import annotations + +import math +import threading +import time +from dataclasses import dataclass +from typing import Callable, Literal + +from tensorrt_llm import logger + + +@dataclass(frozen=True) +class QuiescenceFatalEvent: + """Immutable evidence that an exposed owner missed its retirement deadline.""" + + request_id: int + direction: Literal["send", "receive"] + reason: str + started_at: float + deadline: float + expired_at: float + + +class RetirementDeadline: + """Arbitrate one session's physical claims against a non-resettable deadline.""" + + def __init__( + self, + controller: RetirementWatchdog, + request_id: int, + direction: Literal["send", "receive"], + timeout_s: float, + ) -> None: + """Initialize an unpublished owner without starting either clock. + + Args: + controller: Shared, backend-free watchdog and admission arbiter. + request_id: Request identity within this worker. + direction: Whether this session owns source or destination memory. + timeout_s: Request timeout and subsequent quiescence grace duration. + """ + if not math.isfinite(timeout_s) or timeout_s <= 0: + raise ValueError("retirement requires a finite positive transfer timeout") + self.controller = controller + self.request_id = request_id + self.direction = direction + self.timeout_s = timeout_s + self._request_deadline: float | None = None + self._drain_started: float | None = None + self._reason = "" + self._claims: dict[int, object] = {} + self._pieces_complete = False + self._closed = False + self._timeout_outcome: Callable[[], None] | None = None + + def bind_timeout_outcome(self, callback: Callable[[], None]) -> None: + """Bind the session's metadata-only logical timeout transition. + + Args: + callback: Updates only logical outcome metadata under the shared lock. + It must not call a backend, acquire a session lock, or perform I/O. + """ + with self.lock: + self._timeout_outcome = callback + + def check(self) -> None: + """Commit any elapsed timeout before a logical or physical transition.""" + with self.lock: + self._check_locked(self.controller.clock()) + + @property + def lock(self) -> threading.RLock: + """Return the arbiter shared by exposure, physical settlement and expiry.""" + return self.controller.lock + + def expose(self, *owners: object) -> bool: + """Atomically retain owners before publication or submission can escape. + + Args: + owners: Physical operations or receive owners to keep alive together. + + Returns: + Whether this operation may cross the publication/submission boundary. + """ + with self.lock: + now = self.controller.clock() + self._check_locked(now) + if ( + self._closed + or self._pieces_complete + or self._drain_started is not None + or self.controller.fatal is not None + ): + return False + if self._request_deadline is None: + self._request_deadline = now + self.timeout_s + self._claims.update((id(owner), owner) for owner in owners) + self.controller.wake.set() + return True + + def complete_pieces(self) -> None: + """Seal successful delivery of every expected piece, including required AUX. + + Physical claims remain independently monitored until safe evidence settles + them; successful logical delivery alone cannot disable the deadline. + """ + with self.lock: + self._check_locked(self.controller.clock()) + self._pieces_complete = True + + @property + def is_complete(self) -> bool: + """Return successful whole-session closure, not merely an idle claim set.""" + with self.lock: + self._check_locked(self.controller.clock()) + return ( + self._pieces_complete + and not self._claims + and self._drain_started is None + and self.controller.fatal is None + ) + + def request_drain(self, reason: str) -> None: + """Record the first terminal trigger; later triggers cannot extend grace. + + Args: + reason: Diagnostic explanation of cancellation, failure or shutdown. + """ + with self.lock: + now = self.controller.clock() + self._check_locked(now) + if self._closed or self._drain_started is not None: + return + self._drain_started = now + self._reason = reason + self.controller.wake.set() + + def settle(self, owner: object) -> bool: + """Authorize retirement only if safe evidence wins before fatal expiry. + + Args: + owner: The exact owner previously exposed by this session. + + Returns: + Whether the caller may release that owner's physical roots. + """ + with self.lock: + self._check_locked(self.controller.clock()) + if self.controller.fatal is not None: + return False + self._claims.pop(id(owner), None) + return True + + def retain_unproven(self, owner: object, reason: str) -> None: + """Restore disputed evidence to deadline tracking without reopening admission. + + Args: + owner: An existing physical owner whose safe evidence is now disputed. + reason: Diagnostic explanation for the first ambiguity trigger. + """ + with self.lock: + if self._closed: + return + # A successfully completed session has no running request clock. Start grace + # from this new ambiguity, but preserve any earlier drain trigger. + self.request_drain(reason) + self._claims[id(owner)] = owner + self._check_locked(self.controller.clock()) + self.controller.wake.set() + + def can_retire(self) -> bool: + """Return whether no physical claim remains and fatal expiry has not won.""" + with self.lock: + self._check_locked(self.controller.clock()) + return not self._claims and self.controller.fatal is None + + def close(self) -> bool: + """Idempotently remove a safely retired session from watchdog tracking.""" + with self.lock: + if not self.can_retire(): + return False + self._closed = True + self.controller._owners.discard(self) + return True + + def _check_locked(self, now: float) -> None: + """Latch expiry using metadata only while holding the shared arbiter. + + Args: + now: Monotonic clock sample used for the expiry decision. + """ + if self._closed or self.controller.fatal is not None: + return + session_complete = self._pieces_complete and not self._claims + if ( + not session_complete + and self._drain_started is None + and self._request_deadline is not None + ): + if now >= self._request_deadline: + self._drain_started = self._request_deadline + self._reason = "transfer timeout" + if self._timeout_outcome is not None: + self._timeout_outcome() + if self._claims and self._drain_started is not None: + deadline = self._drain_started + self.timeout_s + if now >= deadline: + self.controller.fatal = QuiescenceFatalEvent( + self.request_id, + self.direction, + self._reason, + self._drain_started, + deadline, + now, + ) + self.controller.wake.set() + + +class RetirementWatchdog: + """Progress deadlines independently from backend waits, queries and executor polling.""" + + def __init__( + self, + callback: Callable[[QuiescenceFatalEvent], None], + *, + clock: Callable[[], float] = time.monotonic, + ) -> None: + """Create a stopped watchdog with a qualified containment callback. + + Args: + callback: Fatal notification; must not perform ordinary resource cleanup. + clock: Monotonic clock, injectable for deterministic deadline tests. + """ + self.lock = threading.RLock() + self.clock = clock + self.wake = threading.Event() + self.fatal: QuiescenceFatalEvent | None = None + self._callback = callback + self._notified = False + self._owners: set[RetirementDeadline] = set() + self._stopped = threading.Event() + self._admission_closed = False + self._thread: threading.Thread | None = None + + def create_owner( + self, request_id: int, direction: Literal["send", "receive"], timeout_s: float + ) -> RetirementDeadline: + """Register one session before it becomes visible to transfer workers. + + Args: + request_id: Session's request identity. + direction: Source or destination ownership role. + timeout_s: Finite positive request and quiescence timeout. + + Returns: + The session's ownership/deadline arbiter. + + Raises: + RuntimeError: Admission is closed following fatal expiry or shutdown. + """ + with self.lock: + if self.fatal is not None or self._admission_closed or self._stopped.is_set(): + raise RuntimeError("KV retirement admission is closed") + owner = RetirementDeadline(self, request_id, direction, timeout_s) + self._owners.add(owner) + return owner + + def start(self) -> None: + """Start the single backend-free progress thread idempotently.""" + with self.lock: + if self._thread is not None: + return + self._thread = threading.Thread(target=self._run, name="kv-retirement", daemon=True) + self._thread.start() + + def progress(self) -> None: + """Evaluate all clocks, then deliver at most one fatal event outside locks.""" + notify = None + with self.lock: + now = self.clock() + for owner in self._owners: + owner._check_locked(now) + if self.fatal is not None and not self._notified: + self._notified = True + notify = self.fatal + if notify is not None: + try: + self._callback(notify) + except Exception as error: + # The containment boundary is external. Failure cannot erase its + # fatal decision or destroy the owners that still retain memory. + logger.error(f"KV retirement containment callback failed; owners retained: {error}") + + def request_shutdown(self) -> None: + """Start drain on every tracked owner before any resource teardown.""" + with self.lock: + self._admission_closed = True + for owner in self._owners: + owner.request_drain("shutdown") + + @property + def admission_closed(self) -> bool: + """Return the irreversible shutdown gate, without entering any session lock.""" + with self.lock: + return self._admission_closed + + def stop(self) -> None: + """Stop progress only after every owner has safely closed. + + Raises: + RuntimeError: Ownership remains tracked, including fatal quarantine. + """ + with self.lock: + if self._owners or self.fatal is not None: + raise RuntimeError("cannot stop KV retirement with retained owners") + self._stopped.set() + self.wake.set() + if self._thread is not None: + self._thread.join(timeout=1) + + def require_retired(self) -> None: + """Refuse ordinary teardown while any session or fatal quarantine is retained. + + Raises: + RuntimeError: A session still owns resources or fatal expiry won. + """ + with self.lock: + if self._owners or self.fatal is not None: + raise RuntimeError("KV retirement still retains resources; teardown refused") + + def _run(self) -> None: + """Poll metadata without acquiring a backend, CUDA, or collective lock.""" + while not self._stopped.is_set(): + self.wake.wait(0.01) + self.wake.clear() + self.progress() + if self.fatal is not None: + # The bound callback roots the worker, sessions, pools and agent. + # If containment returns/raises, retain those roots until process + # teardown; stop() deliberately cannot release a fatal watchdog. + self._stopped.wait() + return diff --git a/tensorrt_llm/_torch/disaggregation/native/retirement.py b/tensorrt_llm/_torch/disaggregation/native/retirement.py index 455376d42905..a282cddce4bd 100644 --- a/tensorrt_llm/_torch/disaggregation/native/retirement.py +++ b/tensorrt_llm/_torch/disaggregation/native/retirement.py @@ -1,346 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Deadline arbitration that never queries or releases a transfer backend.""" - -from __future__ import annotations - -import math -import threading -import time -from dataclasses import dataclass -from typing import Callable, Literal - -from tensorrt_llm import logger - - -@dataclass(frozen=True) -class QuiescenceFatalEvent: - """Immutable evidence that an exposed owner missed its retirement deadline.""" - - request_id: int - direction: Literal["send", "receive"] - reason: str - started_at: float - deadline: float - expired_at: float - - -class RetirementDeadline: - """Arbitrate one session's physical claims against a non-resettable deadline.""" - - def __init__( - self, - controller: RetirementWatchdog, - request_id: int, - direction: Literal["send", "receive"], - timeout_s: float, - ) -> None: - """Initialize an unpublished owner without starting either clock. - - Args: - controller: Shared, backend-free watchdog and admission arbiter. - request_id: Request identity within this worker. - direction: Whether this session owns source or destination memory. - timeout_s: Request timeout and subsequent quiescence grace duration. - """ - if not math.isfinite(timeout_s) or timeout_s <= 0: - raise ValueError("retirement requires a finite positive transfer timeout") - self.controller = controller - self.request_id = request_id - self.direction = direction - self.timeout_s = timeout_s - self._request_deadline: float | None = None - self._drain_started: float | None = None - self._reason = "" - self._claims: dict[int, object] = {} - self._pieces_complete = False - self._closed = False - self._timeout_outcome: Callable[[], None] | None = None - - def bind_timeout_outcome(self, callback: Callable[[], None]) -> None: - """Bind the session's metadata-only logical timeout transition. - - Args: - callback: Updates only logical outcome metadata under the shared lock. - It must not call a backend, acquire a session lock, or perform I/O. - """ - with self.lock: - self._timeout_outcome = callback - - def check(self) -> None: - """Commit any elapsed timeout before a logical or physical transition.""" - with self.lock: - self._check_locked(self.controller.clock()) - - @property - def lock(self) -> threading.RLock: - """Return the arbiter shared by exposure, physical settlement and expiry.""" - return self.controller.lock - - def expose(self, *owners: object) -> bool: - """Atomically retain owners before publication or submission can escape. - - Args: - owners: Physical operations or receive owners to keep alive together. - - Returns: - Whether this operation may cross the publication/submission boundary. - """ - with self.lock: - now = self.controller.clock() - self._check_locked(now) - if ( - self._closed - or self._pieces_complete - or self._drain_started is not None - or self.controller.fatal is not None - ): - return False - if self._request_deadline is None: - self._request_deadline = now + self.timeout_s - self._claims.update((id(owner), owner) for owner in owners) - self.controller.wake.set() - return True - - def complete_pieces(self) -> None: - """Seal successful delivery of every expected piece, including required AUX. - - Physical claims remain independently monitored until safe evidence settles - them; successful logical delivery alone cannot disable the deadline. - """ - with self.lock: - self._check_locked(self.controller.clock()) - self._pieces_complete = True - - @property - def is_complete(self) -> bool: - """Return successful whole-session closure, not merely an idle claim set.""" - with self.lock: - self._check_locked(self.controller.clock()) - return ( - self._pieces_complete - and not self._claims - and self._drain_started is None - and self.controller.fatal is None - ) - - def request_drain(self, reason: str) -> None: - """Record the first terminal trigger; later triggers cannot extend grace. - - Args: - reason: Diagnostic explanation of cancellation, failure or shutdown. - """ - with self.lock: - now = self.controller.clock() - self._check_locked(now) - if self._closed or self._drain_started is not None: - return - self._drain_started = now - self._reason = reason - self.controller.wake.set() - - def settle(self, owner: object) -> bool: - """Authorize retirement only if safe evidence wins before fatal expiry. - - Args: - owner: The exact owner previously exposed by this session. - - Returns: - Whether the caller may release that owner's physical roots. - """ - with self.lock: - self._check_locked(self.controller.clock()) - if self.controller.fatal is not None: - return False - self._claims.pop(id(owner), None) - return True - - def retain_unproven(self, owner: object, reason: str) -> None: - """Restore disputed evidence to deadline tracking without reopening admission. - - Args: - owner: An existing physical owner whose safe evidence is now disputed. - reason: Diagnostic explanation for the first ambiguity trigger. - """ - with self.lock: - if self._closed: - return - # A successfully completed session has no running request clock. Start grace - # from this new ambiguity, but preserve any earlier drain trigger. - self.request_drain(reason) - self._claims[id(owner)] = owner - self._check_locked(self.controller.clock()) - self.controller.wake.set() - - def can_retire(self) -> bool: - """Return whether no physical claim remains and fatal expiry has not won.""" - with self.lock: - self._check_locked(self.controller.clock()) - return not self._claims and self.controller.fatal is None - - def close(self) -> bool: - """Idempotently remove a safely retired session from watchdog tracking.""" - with self.lock: - if not self.can_retire(): - return False - self._closed = True - self.controller._owners.discard(self) - return True - - def _check_locked(self, now: float) -> None: - """Latch expiry using metadata only while holding the shared arbiter. - - Args: - now: Monotonic clock sample used for the expiry decision. - """ - if self._closed or self.controller.fatal is not None: - return - session_complete = self._pieces_complete and not self._claims - if ( - not session_complete - and self._drain_started is None - and self._request_deadline is not None - ): - if now >= self._request_deadline: - self._drain_started = self._request_deadline - self._reason = "transfer timeout" - if self._timeout_outcome is not None: - self._timeout_outcome() - if self._claims and self._drain_started is not None: - deadline = self._drain_started + self.timeout_s - if now >= deadline: - self.controller.fatal = QuiescenceFatalEvent( - self.request_id, - self.direction, - self._reason, - self._drain_started, - deadline, - now, - ) - self.controller.wake.set() - - -class RetirementWatchdog: - """Progress deadlines independently from backend waits, queries and executor polling.""" - - def __init__( - self, - callback: Callable[[QuiescenceFatalEvent], None], - *, - clock: Callable[[], float] = time.monotonic, - ) -> None: - """Create a stopped watchdog with a qualified containment callback. - - Args: - callback: Fatal notification; must not perform ordinary resource cleanup. - clock: Monotonic clock, injectable for deterministic deadline tests. - """ - self.lock = threading.RLock() - self.clock = clock - self.wake = threading.Event() - self.fatal: QuiescenceFatalEvent | None = None - self._callback = callback - self._notified = False - self._owners: set[RetirementDeadline] = set() - self._stopped = threading.Event() - self._admission_closed = False - self._thread: threading.Thread | None = None - - def create_owner( - self, request_id: int, direction: Literal["send", "receive"], timeout_s: float - ) -> RetirementDeadline: - """Register one session before it becomes visible to transfer workers. - - Args: - request_id: Session's request identity. - direction: Source or destination ownership role. - timeout_s: Finite positive request and quiescence timeout. - - Returns: - The session's ownership/deadline arbiter. - - Raises: - RuntimeError: Admission is closed following fatal expiry or shutdown. - """ - with self.lock: - if self.fatal is not None or self._admission_closed or self._stopped.is_set(): - raise RuntimeError("KV retirement admission is closed") - owner = RetirementDeadline(self, request_id, direction, timeout_s) - self._owners.add(owner) - return owner - - def start(self) -> None: - """Start the single backend-free progress thread idempotently.""" - with self.lock: - if self._thread is not None: - return - self._thread = threading.Thread(target=self._run, name="kv-retirement", daemon=True) - self._thread.start() - - def progress(self) -> None: - """Evaluate all clocks, then deliver at most one fatal event outside locks.""" - notify = None - with self.lock: - now = self.clock() - for owner in self._owners: - owner._check_locked(now) - if self.fatal is not None and not self._notified: - self._notified = True - notify = self.fatal - if notify is not None: - try: - self._callback(notify) - except Exception as error: - # The containment boundary is external. Failure cannot erase its - # fatal decision or destroy the owners that still retain memory. - logger.error(f"KV retirement containment callback failed; owners retained: {error}") - - def request_shutdown(self) -> None: - """Start drain on every tracked owner before any resource teardown.""" - with self.lock: - self._admission_closed = True - for owner in self._owners: - owner.request_drain("shutdown") - - @property - def admission_closed(self) -> bool: - """Return the irreversible shutdown gate, without entering any session lock.""" - with self.lock: - return self._admission_closed - - def stop(self) -> None: - """Stop progress only after every owner has safely closed. - - Raises: - RuntimeError: Ownership remains tracked, including fatal quarantine. - """ - with self.lock: - if self._owners or self.fatal is not None: - raise RuntimeError("cannot stop KV retirement with retained owners") - self._stopped.set() - self.wake.set() - if self._thread is not None: - self._thread.join(timeout=1) - - def require_retired(self) -> None: - """Refuse ordinary teardown while any session or fatal quarantine is retained. - - Raises: - RuntimeError: A session still owns resources or fatal expiry won. - """ - with self.lock: - if self._owners or self.fatal is not None: - raise RuntimeError("KV retirement still retains resources; teardown refused") - - def _run(self) -> None: - """Poll metadata without acquiring a backend, CUDA, or collective lock.""" - while not self._stopped.is_set(): - self.wake.wait(0.01) - self.wake.clear() - self.progress() - if self.fatal is not None: - # The bound callback roots the worker, sessions, pools and agent. - # If containment returns/raises, retain those roots until process - # teardown; stop() deliberately cannot release a fatal watchdog. - self._stopped.wait() - return +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Compatibility imports for the shared lifecycle deadline controller.""" + +from ..lifecycle.retirement import QuiescenceFatalEvent, RetirementDeadline, RetirementWatchdog + +__all__ = ["QuiescenceFatalEvent", "RetirementDeadline", "RetirementWatchdog"] diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 13612283b77d..df3d66b14294 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -56,6 +56,21 @@ TxSessionBase, WaitResult, ) +from tensorrt_llm._torch.disaggregation.lifecycle.ownership import ( + PhysicalOperationState as _PhysicalOperationState, # noqa: F401 -- native diagnostic alias +) +from tensorrt_llm._torch.disaggregation.lifecycle.ownership import ( + ReceiveOperationOwner as _ReceiveOperationOwner, +) +from tensorrt_llm._torch.disaggregation.lifecycle.ownership import SendOperationOwner +from tensorrt_llm._torch.disaggregation.lifecycle.ownership import ( + TransferNotSubmittedError as _TransferNotSubmittedError, +) +from tensorrt_llm._torch.disaggregation.lifecycle.retirement import ( + QuiescenceFatalEvent, + RetirementDeadline, + RetirementWatchdog, +) from tensorrt_llm._torch.disaggregation.native.auxiliary import ( AuxBuffer, build_aux_transfer_layout, @@ -70,11 +85,6 @@ from tensorrt_llm._torch.disaggregation.native.peer import PeerOverlap, PeerRegistrar from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfTimer, perf_log_manager from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo -from tensorrt_llm._torch.disaggregation.native.retirement import ( - QuiescenceFatalEvent, - RetirementDeadline, - RetirementWatchdog, -) from tensorrt_llm._torch.disaggregation.native.utils import get_local_ip from tensorrt_llm._torch.disaggregation.nixl.agent import NixlTransferAgent from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 @@ -86,6 +96,8 @@ from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle if TYPE_CHECKING: + from tensorrt_llm._torch.disaggregation.base.agent import TransferStatus + from .bounce import Config AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType @@ -103,10 +115,6 @@ _FALLBACK_TX_OVERALL_TIMEOUT_S = 60.0 -class _TransferNotSubmittedError(RuntimeError): - """The sender rejected an operation before the backend could access memory.""" - - @dataclass class RecvReqInfo: sender_req_id: int @@ -329,251 +337,6 @@ 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.""" - - def __init__(self, retirement: Optional[RetirementDeadline] = None) -> None: - self._retirement = retirement - self._lock = threading.Lock() if retirement is None else retirement.lock - self._publication_pending = False - self._cancelled_unpublished = False - self._expected_writers: Optional[int] = None - self._writer_cohort: Optional[frozenset[int]] = None - self._writer_candidates: Optional[frozenset[int]] = None - self._published_writers: Optional[frozenset[int]] = None - self._quiesced_sessions: set[int] = set() - self._writer_results: dict[int, bool] = {} - self._in_doubt_writers: set[int] = set() - self._settled_writers: set[int] = set() - self._publication_failed = False - self._local_completion_pending = False - self._invalid_evidence = False - - def _settle_if_drained_locked(self) -> None: - """Remove this claim only after the complete cohort and local work settle.""" - if self._retirement is not None and self._resources_drained_locked(): - self._retirement.settle(self) - - def _invalidate_evidence_locked(self) -> None: - """Quarantine conflicting proof, including proof disputed after settlement.""" - self._invalid_evidence = True - if self._retirement is not None: - self._retirement.retain_unproven(self, "invalid receive ownership evidence") - - def begin_publication(self) -> None: - with self._lock: - if self._publication_pending or self._expected_writers is not None: - raise RuntimeError("destination publication was already started") - self._publication_pending = True - - def seal_writer_cohort( - self, - expected_writers: int, - writer_cohort: Optional[set[int]] = None, - *, - published_writers: Optional[set[int]] = None, - ) -> None: - if expected_writers < 0: - raise ValueError(f"expected_writers must be non-negative, got {expected_writers}") - cohort = None if writer_cohort is None else frozenset(writer_cohort) - candidates = cohort if published_writers is None else frozenset(published_writers) - if cohort is not None and len(cohort) != expected_writers: - raise ValueError( - f"writer cohort has {len(cohort)} member(s), expected {expected_writers}" - ) - if candidates is not None and ( - len(candidates) < expected_writers or (cohort is not None and not cohort <= candidates) - ): - raise ValueError("published candidates must cover every eligible writer") - with self._lock: - if self._expected_writers is not None: - if ( - self._expected_writers != expected_writers - or self._writer_cohort != cohort - or self._writer_candidates != candidates - ): - raise RuntimeError("writer cohort was already sealed differently") - return - self._expected_writers = expected_writers - self._writer_cohort = cohort - self._writer_candidates = candidates - self._published_writers = candidates - - def finish_publication(self) -> None: - """Record that every authorized REQUEST_DATA message was sent.""" - with self._lock: - self._publication_pending = False - self._settle_if_drained_locked() - - def abort_publication(self, published_writers: set[int]) -> None: - """Close a failed fan-out around the writers whose sends succeeded.""" - with self._lock: - published = frozenset(published_writers) - if self._writer_cohort is not None and not published.issubset(self._writer_cohort): - self._invalidate_evidence_locked() - raise RuntimeError("publication recorded a writer outside the sealed cohort") - if not (self._writer_results.keys() | self._in_doubt_writers) <= published: - self._invalidate_evidence_locked() - raise RuntimeError("terminal evidence came from an unpublished writer") - self._expected_writers = len(published) - self._writer_cohort = frozenset(published) - self._published_writers = published - self._publication_failed = True - self._cancelled_unpublished = not published - self._publication_pending = False - self._settle_if_drained_locked() - - def cancel_unpublished(self) -> bool: - """Close a publication that did not authorize a remote writer.""" - with self._lock: - if self._invalid_evidence: - return False - if self._expected_writers is not None: - return self._cancelled_unpublished - self._cancelled_unpublished = True - self._expected_writers = 0 - self._writer_cohort = frozenset() - self._published_writers = frozenset() - self._publication_pending = False - self._settle_if_drained_locked() - return True - - def record_session_quiesced(self, peer_rank: int) -> None: - """Record no-future-access proof without inventing a per-piece result.""" - with self._lock: - if self._writer_candidates is None or peer_rank not in self._writer_candidates: - self._invalidate_evidence_locked() - raise RuntimeError(f"session acknowledgment from unknown writer {peer_rank}") - if self._published_writers is not None and peer_rank in self._published_writers: - self._quiesced_sessions.add(peer_rank) - self._settle_if_drained_locked() - - def record_writer_in_doubt(self, peer_rank: int) -> bool: - """Retain ownership after a writer reports no safe terminal evidence.""" - with self._lock: - if self._expected_writers is None: - self._invalidate_evidence_locked() - raise RuntimeError( - f"writer {peer_rank} reported ambiguous evidence before publication" - ) - if self._writer_cohort is not None and peer_rank not in self._writer_cohort: - self._invalidate_evidence_locked() - raise RuntimeError(f"writer {peer_rank} is outside the sealed cohort") - if peer_rank in self._in_doubt_writers or peer_rank in self._settled_writers: - return False - self._in_doubt_writers.add(peer_rank) - if self._retirement is not None: - self._retirement.retain_unproven(self, "receive ownership is in doubt") - return True - - def record_writer_settlement(self, peer_rank: int) -> bool: - """Accept physical DONE for a writer that previously reported IN_DOUBT. - - Ordinary FAILED is not this proof: it may describe a later, unsubmitted - chunk while the earlier ambiguous write is still touching the destination. - """ - with self._lock: - if self._expected_writers is None or ( - self._writer_cohort is not None and peer_rank not in self._writer_cohort - ): - self._invalidate_evidence_locked() - raise RuntimeError(f"writer {peer_rank} settled outside the published cohort") - if peer_rank in self._settled_writers: - return False - if peer_rank not in self._in_doubt_writers: - self._invalidate_evidence_locked() - raise RuntimeError(f"writer {peer_rank} settled without prior ambiguous evidence") - if self._writer_results.get(peer_rank) is True: - self._invalidate_evidence_locked() - raise RuntimeError(f"writer {peer_rank} settled after contradictory success") - self._writer_results[peer_rank] = False - self._in_doubt_writers.remove(peer_rank) - self._settled_writers.add(peer_rank) - self._settle_if_drained_locked() - return True - - def record_writer_result( - self, - peer_rank: int, - succeeded: bool, - *, - wait_for_local_completion: bool, - ) -> tuple[bool, bool]: - """Record one writer and return ``(accepted, all_succeeded)``.""" - with self._lock: - if self._expected_writers is None: - self._invalidate_evidence_locked() - raise RuntimeError( - f"writer {peer_rank} reported terminal evidence before publication" - ) - if self._writer_cohort is not None and peer_rank not in self._writer_cohort: - self._invalidate_evidence_locked() - raise RuntimeError(f"writer {peer_rank} is outside the sealed cohort") - if peer_rank in self._in_doubt_writers: - if succeeded: - self._invalidate_evidence_locked() - raise RuntimeError(f"writer {peer_rank} reported success while in doubt") - return False, False - previous = self._writer_results.get(peer_rank) - if previous is not None: - if previous != succeeded: - self._invalidate_evidence_locked() - raise RuntimeError( - f"writer {peer_rank} reported contradictory terminal evidence" - ) - return False, False - if len(self._writer_results) >= self._expected_writers: - return False, False - self._writer_results[peer_rank] = succeeded - all_reported = len(self._writer_results) == self._expected_writers - all_succeeded = ( - all_reported and not self._publication_failed and all(self._writer_results.values()) - ) - if all_succeeded and wait_for_local_completion: - self._local_completion_pending = True - self._settle_if_drained_locked() - return True, all_succeeded - - def finish_local_completion(self) -> None: - with self._lock: - self._local_completion_pending = False - self._settle_if_drained_locked() - - @property - def all_writers_reported(self) -> bool: - """Whether every writer of the sealed cohort has reported a terminal result.""" - with self._lock: - return ( - self._expected_writers is not None - and len(self._writer_results) == self._expected_writers - and not self._in_doubt_writers - ) - - @property - def resources_drained(self) -> bool: - with self._lock: - return self._resources_drained_locked() and ( - self._retirement is None or self._retirement.can_retire() - ) - - def _resources_drained_locked(self) -> bool: - """Check physical evidence without consulting the session's other owners.""" - return ( - self._expected_writers is not None - and ( - len(self._writer_results) == self._expected_writers - or ( - self._published_writers is not None - and self._published_writers <= self._quiesced_sessions - ) - ) - and not self._publication_pending - and not self._local_completion_pending - and not self._in_doubt_writers - and not self._invalid_evidence - ) - - class AgentResult(Enum): SUCCESS = "SUCCESS" FAILED = "FAILED" @@ -630,44 +393,6 @@ def _make_aux_result_msg(instance_rank: int, unique_rid: int, result: AgentResul ] -class _PhysicalOperationState(Enum): - """Sender-side evidence for one peer's access to a task's source memory. - - The legal forward paths are:: - - ADMITTED -> NOT_SUBMITTED - ADMITTED -> SUBMITTING -> SUBMITTED -> BACKEND_DONE - SUBMITTING -> IN_DOUBT - SUBMITTED -> IN_DOUBT -> BACKEND_DONE (same retained status reports DONE) - - Only NOT_SUBMITTED and BACKEND_DONE prove that the operation can no longer - access the source. An IN_DOUBT operation without a retained backend status - cannot retire. Repeating a terminal transition is idempotent, and repeating - IN_DOUBT preserves the retained backend evidence. - """ - - ADMITTED = "ADMITTED" - SUBMITTING = "SUBMITTING" - SUBMITTED = "SUBMITTED" - NOT_SUBMITTED = "NOT_SUBMITTED" - BACKEND_DONE = "BACKEND_DONE" - IN_DOUBT = "IN_DOUBT" - - -_DRAINED_PHYSICAL_OPERATION_STATES = frozenset( - (_PhysicalOperationState.NOT_SUBMITTED, _PhysicalOperationState.BACKEND_DONE) -) - - -@dataclass -class _PhysicalOperation: - """State plus strong backend roots retained until physical quiescence.""" - - state: _PhysicalOperationState - request: Optional[TransferRequest] = None - status: Optional[object] = None - - @dataclass class _PendingSettlement: write_meta: WriteMeta @@ -687,8 +412,9 @@ def __init__(self, params: DisaggregatedParams): self._params = params self._unique_rid: Optional[int] = params.disagg_request_id self._perf_timer = PerfTimer() if perf_log_manager.enabled else None - self._physical_lock = threading.Lock() - self._physical_operations: dict[int, _PhysicalOperation] = {} + self._physical_owner = SendOperationOwner() + # Preserve the native diagnostic view of the shared owner's exact roots. + self._physical_operations = self._physical_owner._physical_operations def bind_logical_outcomes(self, outcomes: _LogicalOutcomes) -> None: """Bind session outcome and deadline before exposing the task to workers. @@ -697,8 +423,7 @@ def bind_logical_outcomes(self, outcomes: _LogicalOutcomes) -> None: outcomes: The session's stable logical-result arbiter. """ super().bind_logical_outcomes(outcomes) - if self._retirement is not None: - self._physical_lock = self._retirement.lock + self._physical_owner.bind_retirement(self._retirement) def fail(self, exc: Exception) -> None: self._logical_outcomes.fail(exc) @@ -720,142 +445,41 @@ def is_done(self) -> bool: return self._event.is_set() def begin_physical_operation(self, peer_rank: int) -> bool: - with self._physical_lock: - if peer_rank in self._physical_operations: - return False - self._physical_operations[peer_rank] = _PhysicalOperation( - _PhysicalOperationState.ADMITTED - ) - return True + """Admit a participant once in the shared physical owner.""" + return self._physical_owner.begin_physical_operation(peer_rank) - def _require_physical_operation_locked( - self, - peer_rank: int, - expected_states: tuple[_PhysicalOperationState, ...], - ) -> _PhysicalOperation: - operation = self._physical_operations.get(peer_rank) - if operation is None: - raise RuntimeError(f"physical operation {peer_rank} was not admitted") - if operation.state not in expected_states: - expected = ", ".join(state.value for state in expected_states) - raise RuntimeError( - f"physical operation {peer_rank} is {operation.state.value}, expected {expected}" - ) - return operation + def begin_backend_submission(self, peer_rank: int, request: TransferRequest) -> None: + """Root the request before the native backend submission escapes.""" + self._physical_owner.begin_backend_submission(peer_rank, request) - def begin_backend_submission( - self, - peer_rank: int, - request: TransferRequest, - ) -> None: - with self._physical_lock: - operation = self._require_physical_operation_locked( - peer_rank, (_PhysicalOperationState.ADMITTED,) - ) - if self._retirement is not None and not self._retirement.expose(operation): - operation.state = _PhysicalOperationState.NOT_SUBMITTED - raise _TransferNotSubmittedError("source retirement admission is closed") - operation.request = request - operation.state = _PhysicalOperationState.SUBMITTING - - def record_backend_submission(self, peer_rank: int, status: object) -> None: - with self._physical_lock: - operation = self._require_physical_operation_locked( - peer_rank, (_PhysicalOperationState.SUBMITTING,) - ) - operation.status = status - operation.state = _PhysicalOperationState.SUBMITTED + def record_backend_submission(self, peer_rank: int, status: TransferStatus) -> None: + """Retain the exact status returned by the native backend.""" + self._physical_owner.record_backend_submission(peer_rank, status) def mark_physical_operation_in_doubt(self, peer_rank: int) -> None: - with self._physical_lock: - operation = self._require_physical_operation_locked( - peer_rank, - ( - _PhysicalOperationState.SUBMITTING, - _PhysicalOperationState.SUBMITTED, - _PhysicalOperationState.IN_DOUBT, - ), - ) - operation.state = _PhysicalOperationState.IN_DOUBT - if self._retirement is not None: - self._retirement.request_drain("backend quiescence unproven") + """Retain ambiguous access independently of the logical outcome.""" + self._physical_owner.mark_physical_operation_in_doubt(peer_rank) def retire_unsubmitted_physical_operation(self, peer_rank: int) -> None: - with self._physical_lock: - operation = self._require_physical_operation_locked( - peer_rank, - ( - _PhysicalOperationState.ADMITTED, - _PhysicalOperationState.NOT_SUBMITTED, - ), - ) - operation.state = _PhysicalOperationState.NOT_SUBMITTED + """Record that a participant never submitted a backend operation.""" + self._physical_owner.retire_unsubmitted_physical_operation(peer_rank) def retire_backend_done_physical_operation(self, peer_rank: int) -> bool: - with self._physical_lock: - operation = self._require_physical_operation_locked( - peer_rank, - ( - _PhysicalOperationState.SUBMITTED, - _PhysicalOperationState.BACKEND_DONE, - ), - ) - if operation.state is _PhysicalOperationState.BACKEND_DONE: - return True - if self._retirement is not None and not self._retirement.settle(operation): - operation.state = _PhysicalOperationState.IN_DOUBT - return False - operation.request = None - operation.status = None - operation.state = _PhysicalOperationState.BACKEND_DONE - return True + """Settle backend-defined completion through the common deadline arbiter.""" + return self._physical_owner.retire_backend_done_physical_operation(peer_rank) def poll_in_doubt_physical_operation(self, peer_rank: int) -> bool: - """Retire once, only after a fresh DONE query on the retained status. - - Polling does not change the task's logical outcome. Keep strong local - roots across the query, and reject a result if its operation changed. - """ - with self._physical_lock: - operation = self._physical_operations.get(peer_rank) - if operation is None or operation.state is not _PhysicalOperationState.IN_DOUBT: - return False - request, status = operation.request, operation.status - if status is None: - return False - try: - completed = status.is_completed() - except Exception: - # A backend query failure is not evidence that its accessors stopped. - return False - if completed is not True: - return False - with self._physical_lock: - if ( - self._physical_operations.get(peer_rank) is not operation - or operation.state is not _PhysicalOperationState.IN_DOUBT - or operation.request is not request - or operation.status is not status - ): - return False - if self._retirement is not None and not self._retirement.settle(operation): - return False - operation.request = None - operation.status = None - operation.state = _PhysicalOperationState.BACKEND_DONE - return True + """Query and settle only the exact retained ambiguous operation.""" + return self._physical_owner.poll_in_doubt_physical_operation(peer_rank) def has_started_physical_operation(self, peer_rank: int) -> bool: - with self._physical_lock: - return peer_rank in self._physical_operations + """Whether this task has admitted the participant.""" + return self._physical_owner.has_started_physical_operation(peer_rank) @property def resources_drained(self) -> bool: - with self._physical_lock: - return (self._retirement is None or self._retirement.can_retire()) and all( - operation.state in _DRAINED_PHYSICAL_OPERATION_STATES - for operation in self._physical_operations.values() - ) + """Whether operation evidence and the session deadline permit retirement.""" + return self._physical_owner.resources_drained def print_perf_info(self, peer_rank: int, instance_name: str, instance_rank: int): if self._perf_timer is None: diff --git a/tests/unittest/disaggregated/test_lifecycle_binding.py b/tests/unittest/disaggregated/test_lifecycle_binding.py new file mode 100644 index 000000000000..03caa6fcf704 --- /dev/null +++ b/tests/unittest/disaggregated/test_lifecycle_binding.py @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""CPU binding contracts for real, backend-neutral lifecycle primitives. + +This adapter-shaped fixture is not an RI-02 implementation or profile +qualification. Only manager holds, backend status, and already committed logical +outcomes are doubles; admission, evidence tracking, and retirement are production +code. Allocation release remains the caller's responsibility after access ends. +""" + +import gc +import weakref +from dataclasses import dataclass +from unittest.mock import Mock + +import pytest + +from tensorrt_llm._torch.disaggregation.lifecycle.ownership import ( + PhysicalOperationState, + ReceiveOperationOwner, + SendOperationOwner, + TransferNotSubmittedError, +) +from tensorrt_llm._torch.disaggregation.lifecycle.retirement import ( + QuiescenceFatalEvent, + RetirementDeadline, + RetirementWatchdog, +) + +pytestmark = pytest.mark.cpu_only + + +class _Root: + """Opaque, weak-referenceable allocation hold or backend request.""" + + +@dataclass +class _Status: + """Backend-owned completion evidence; no lifecycle transitions live here.""" + + result: object = False + queries: int = 0 + + def is_completed(self) -> object: + """Return the injected backend evidence, or its query failure.""" + self.queries += 1 + if isinstance(self.result, Exception): + raise RuntimeError(str(self.result)) + return self.result + + +class _Binding: + """Compose shared ownership with opaque caller-owned resources and outcomes.""" + + def __init__(self) -> None: + """Create an unpublished source binding without starting a worker thread.""" + self.clock = Mock(return_value=10.0) + self.events: list[QuiescenceFatalEvent] = [] + self.watchdog = RetirementWatchdog(self.events.append, clock=self.clock) + self.retirement = self.watchdog.create_owner(4, "send", 5.0) + self.sender = SendOperationOwner(self.retirement) + self.manager_hold = _Root() + # Fixture-owned committed metadata, not a replacement logical arbiter. + self.logical_outcome = object() + + +@pytest.fixture +def binding() -> _Binding: + """Provide a small non-native caller of the shared implementation.""" + return _Binding() + + +def test_native_compatibility_names_are_shared_classes() -> None: + """Compatibility imports must not fork state or exception identities.""" + from tensorrt_llm._torch.disaggregation.native import retirement, transfer + + assert retirement.RetirementWatchdog is RetirementWatchdog + assert retirement.RetirementDeadline is RetirementDeadline + assert retirement.QuiescenceFatalEvent is QuiescenceFatalEvent + assert transfer._ReceiveOperationOwner is ReceiveOperationOwner + assert transfer._PhysicalOperationState is PhysicalOperationState + assert transfer._TransferNotSubmittedError is TransferNotSubmittedError + assert {state.value for state in PhysicalOperationState} == { + "ADMITTED", + "SUBMITTING", + "SUBMITTED", + "NOT_SUBMITTED", + "BACKEND_DONE", + "IN_DOUBT", + } + + +def test_multi_owner_exposure_waits_for_local_completion(binding: _Binding) -> None: + """KV, AUX, and manager holds survive partial and duplicate remote evidence.""" + binding.retirement.close() + retirement = binding.watchdog.create_owner(5, "receive", 5.0) + kv, aux = ReceiveOperationOwner(retirement), ReceiveOperationOwner(retirement) + assert retirement.expose(kv, aux, binding.manager_hold) + for owner, cohort in ((kv, {7, 8}), (aux, {7})): + owner.begin_publication() + owner.seal_writer_cohort(len(cohort), cohort) + owner.finish_publication() + assert kv.record_writer_result(7, True, wait_for_local_completion=True) == (True, False) + assert kv.record_writer_result(7, True, wait_for_local_completion=True) == (False, False) + assert kv.record_writer_result(8, True, wait_for_local_completion=True) == (True, True) + assert not kv.resources_drained + assert not retirement.can_retire() + assert aux.record_writer_result(7, True, wait_for_local_completion=False) == (True, True) + retirement.complete_pieces() + assert not retirement.is_complete + assert id(kv) in retirement._claims + kv.finish_local_completion() + kv.finish_local_completion() + assert set(retirement._claims) == {id(binding.manager_hold)} + # The caller releases its allocation hold only after both accessors ended. + assert retirement.settle(binding.manager_hold) + assert kv.resources_drained and aux.resources_drained + assert retirement.is_complete + assert retirement.close() and retirement.close() + binding.watchdog.require_retired() + binding.watchdog.stop() + + +@pytest.mark.parametrize("rejected", [False, True]) +def test_never_submitted_operation_retires_idempotently(binding: _Binding, rejected: bool) -> None: + """Skipping submission and closed admission both establish no-access proof.""" + sender = binding.sender + assert sender.begin_physical_operation(7) + assert not sender.begin_physical_operation(7) + if rejected: + binding.retirement.request_drain("cancelled before submission") + with pytest.raises(TransferNotSubmittedError, match="admission is closed"): + sender.begin_backend_submission(7, _Root()) + sender.retire_unsubmitted_physical_operation(7) + sender.retire_unsubmitted_physical_operation(7) + assert sender.has_started_physical_operation(7) + assert sender.resources_drained + assert not binding.retirement._claims + assert binding.retirement.close() and binding.retirement.close() + binding.watchdog.stop() + + +@pytest.mark.parametrize("evidence", [False, 1, RuntimeError("query failed")]) +def test_late_exact_done_retires_without_changing_logical_outcome( + binding: _Binding, evidence: object +) -> None: + """Only literal DONE on the retained status ends ambiguous physical access.""" + sender, status, request = binding.sender, _Status(evidence), _Root() + request_ref, status_ref = weakref.ref(request), weakref.ref(status) + committed = binding.logical_outcome + assert sender.begin_physical_operation(7) + sender.begin_backend_submission(7, request) + sender.record_backend_submission(7, status) + sender.mark_physical_operation_in_doubt(7) + sender.mark_physical_operation_in_doubt(7) + del request + gc.collect() + assert request_ref() is not None + assert not sender.poll_in_doubt_physical_operation(7) + assert not sender.resources_drained + assert binding.logical_outcome is committed + status.result = True + binding.clock.return_value = 14.999 + assert sender.poll_in_doubt_physical_operation(7) + assert sender.resources_drained + assert sender.retire_backend_done_physical_operation(7) + assert not sender.poll_in_doubt_physical_operation(7) + assert status.queries == 2 + assert binding.logical_outcome is committed + del status + gc.collect() + assert request_ref() is None and status_ref() is None + assert binding.retirement.close() and binding.retirement.close() + binding.watchdog.stop() + + +@pytest.mark.parametrize("has_status", [False, True]) +def test_exact_expiry_contains_and_retains_all_roots(binding: _Binding, has_status: bool) -> None: + """Missing handles and exact-boundary DONE both retain roots after fatal expiry.""" + sender, request, status = binding.sender, _Root(), _Status() + request_ref, hold_ref = weakref.ref(request), weakref.ref(binding.manager_hold) + assert binding.retirement.expose(binding.manager_hold) + assert sender.begin_physical_operation(7) + sender.begin_backend_submission(7, request) + if has_status: + sender.record_backend_submission(7, status) + sender.mark_physical_operation_in_doubt(7) + binding.clock.return_value = 14.0 + sender.mark_physical_operation_in_doubt(7) + assert not sender.poll_in_doubt_physical_operation(7) + assert not sender.resources_drained + binding.clock.return_value = 15.0 + status.result = True + assert not sender.poll_in_doubt_physical_operation(7) + binding.watchdog.progress() + binding.watchdog.progress() + event = binding.watchdog.fatal + assert event is not None + assert (event.started_at, event.deadline, event.expired_at) == (10.0, 15.0, 15.0) + assert binding.events == [event] + assert not binding.retirement.settle(binding.manager_hold) + assert not binding.retirement.close() + assert not sender.resources_drained + with pytest.raises(RuntimeError, match="admission is closed"): + binding.watchdog.create_owner(6, "receive", 5.0) + with pytest.raises(RuntimeError, match="teardown refused"): + binding.watchdog.require_retired() + watchdog = binding.watchdog + # Pytest retains its fixture; remove the caller's roots before collecting. + del binding.manager_hold, binding.sender + del request, sender, binding + gc.collect() + assert request_ref() is not None and hold_ref() is not None + with pytest.raises(RuntimeError, match="retained owners"): + watchdog.stop() + + +def test_receive_ambiguity_requires_explicit_settlement(binding: _Binding) -> None: + """An ordinary failed result cannot substitute for access-end evidence.""" + owner = ReceiveOperationOwner(binding.retirement) + assert binding.retirement.expose(owner, binding.manager_hold) + owner.begin_publication() + owner.seal_writer_cohort(1, {7}) + owner.finish_publication() + assert owner.record_writer_in_doubt(7) + assert not owner.record_writer_in_doubt(7) + assert owner.record_writer_result(7, False, wait_for_local_completion=False) == (False, False) + assert not owner.resources_drained + assert owner.record_writer_settlement(7) + assert not owner.record_writer_settlement(7) + assert set(binding.retirement._claims) == {id(binding.manager_hold)} + assert binding.retirement.settle(binding.manager_hold) + assert owner.resources_drained + assert binding.retirement.close() and binding.retirement.close() + binding.watchdog.stop()