diff --git a/tensorrt_llm/_torch/disaggregation/async_consensus.py b/tensorrt_llm/_torch/disaggregation/async_consensus.py new file mode 100644 index 000000000000..11944c3f7185 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/async_consensus.py @@ -0,0 +1,1691 @@ +# 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. + +"""Nonblocking agreement for Python disaggregated KV-transfer state. + +The protocol deliberately separates two state machines: + +* terminal votes are immutable and produce one authoritative outcome after + every participant is locally quiescent; and +* generation-first readiness is withdrawable until a rank grants its prepare + lease; the authoritative PP schedule then activates the same request on + every rank before acknowledged completion retires the in-flight epoch. + +The MPI transport uses a duplicated communicator and fixed-width packets. It +never introduces a collective after construction. +""" + +from __future__ import annotations + +import time +from collections import OrderedDict, deque +from collections.abc import Callable, Iterable, Sequence +from dataclasses import dataclass +from enum import IntEnum +from typing import Protocol + +import numpy as np + +from tensorrt_llm._utils import mpi_comm + +try: + from mpi4py import MPI +except ImportError: + MPI = None + + +PROTOCOL_VERSION = 3 +_PACKET_FIELDS = 7 +_DEFAULT_MAX_MESSAGES_PER_POLL = 256 +_DEFAULT_MAX_COMPLETED_EPOCHS = 65_536 +_DEFAULT_MAX_PENDING_SENDS = 65_536 +_DEFAULT_MAX_OPEN_ROUNDS = 65_536 +_DEFAULT_ROUND_TIMEOUT_S = 600.0 +_DEFAULT_MAX_SEND_TESTS_PER_PROGRESS = 256 + + +class _ConsensusBackpressure(RuntimeError): + """Transient capacity rejection before any MPI send is issued.""" + + +class ConsensusPhase(IntEnum): + READY = 1 + TERMINAL = 2 + + +class ConsensusOutcome(IntEnum): + READY = 1 + COMPLETED = 2 + FAILED = 3 + CANCELLED = 4 + WITHDRAWN = 5 + + +class ConsensusEventKind(IntEnum): + READY_PREPARE = 1 + READY_RELEASE = 2 + READY_COMPLETE = 3 + READY_ABORT = 4 + TERMINAL_COMMIT = 5 + READY_ABORT_FINALIZE = 6 + + +class _MessageKind(IntEnum): + VOTE = 1 + WITHDRAW = 2 + READY_PREPARE = 3 + READY_ACK = 4 + READY_RELEASE = 5 + READY_ACTIVATE_ACK = 6 + READY_ABORT = 7 + TERMINAL_COMMIT = 8 + CLOSE = 9 + CLOSE_ACK = 10 + READY_ABORT_ACK = 11 + READY_ABORT_FINALIZE = 12 + FAIL_STOP = 13 + READY_COMPLETE = 14 + + +class _CoordinatorAction(IntEnum): + VOTES_COMPLETE = 1 + READY_ACKS_COMPLETE = 2 + READY_ABORT_ACKS_COMPLETE = 3 + READY_ABORT_START = 4 + FAIL_STOP = 5 + READY_ACTIVATE_ACKS_COMPLETE = 6 + + +@dataclass(frozen=True) +class ConsensusEvent: + kind: ConsensusEventKind + request_id: int + epoch: int + outcome: ConsensusOutcome + + +@dataclass(frozen=True) +class _Packet: + kind: _MessageKind + phase: ConsensusPhase + request_id: int + epoch: int + outcome: ConsensusOutcome + source: int + + def encode(self) -> np.ndarray: + return np.asarray( + [ + PROTOCOL_VERSION, + int(self.kind), + int(self.phase), + self.request_id, + self.epoch, + int(self.outcome), + self.source, + ], + dtype=np.uint64, + ) + + @classmethod + def decode(cls, fields: np.ndarray) -> "_Packet": + if fields.shape != (_PACKET_FIELDS,): + raise RuntimeError(f"invalid consensus packet shape: {fields.shape}") + if int(fields[0]) != PROTOCOL_VERSION: + raise RuntimeError( + f"unsupported consensus protocol version {int(fields[0])}; " + f"expected {PROTOCOL_VERSION}" + ) + try: + return cls( + kind=_MessageKind(int(fields[1])), + phase=ConsensusPhase(int(fields[2])), + request_id=int(fields[3]), + epoch=int(fields[4]), + outcome=ConsensusOutcome(int(fields[5])), + source=int(fields[6]), + ) + except ValueError as error: + raise RuntimeError(f"invalid consensus packet fields: {fields.tolist()}") from error + + +class ConsensusTransport(Protocol): + rank: int + participants: tuple[int, ...] + + def send(self, packet: _Packet, destination: int) -> None: ... + + def send_many(self, messages: Sequence[tuple[_Packet, int]]) -> None: + """Atomically accept a fan-out or reject it before sending. + + `_ConsensusBackpressure` is the only recoverable rejection. Once any + packet is accepted, an implementation must either accept the complete + batch or terminate the participant group; returning an error after a + strict prefix would make an authoritative decision unsafe to retry. + """ + ... + + def progress(self) -> None: ... + + def receive(self, limit: int) -> list[_Packet]: ... + + @property + def pending_send_count(self) -> int: ... + + def close(self, timeout_s: float) -> None: ... + + +@dataclass +class _PendingSend: + buffer: np.ndarray + request: object + + +@dataclass +class _RoundProgress: + started_at: float + last_progress_at: float + + +class MpiConsensusTransport: + """Fixed-buffer point-to-point transport on a dedicated communicator.""" + + def __init__( + self, + participants: Sequence[int], + *, + max_pending_sends: int = _DEFAULT_MAX_PENDING_SENDS, + max_send_tests_per_progress: int = _DEFAULT_MAX_SEND_TESTS_PER_PROGRESS, + ): + if MPI is None: + raise RuntimeError("mpi4py is required for asynchronous Python consensus") + if max_pending_sends <= 0: + raise ValueError("max_pending_sends must be positive") + if max_send_tests_per_progress <= 0: + raise ValueError("max_send_tests_per_progress must be positive") + participants_tuple = tuple(int(rank) for rank in participants) + minimum_fanout = max(1, len(participants_tuple) - 1) + if max_pending_sends < minimum_fanout: + raise ValueError( + "max_pending_sends must accommodate one participant fan-out: " + f"minimum={minimum_fanout}, configured={max_pending_sends}" + ) + world = mpi_comm() + world_size = world.Get_size() + # The initial qualified topology is TP1/CP1/PP>1, so the PP domain is + # the full worker world. Rejecting subsets avoids communicator-rank + # translation mistakes until attention-DP lanes are qualified. + if participants_tuple != tuple(range(world_size)): + raise RuntimeError( + "asynchronous Python consensus currently requires the participant " + "domain to equal the MPI worker world" + ) + self.participants = participants_tuple + comm = world.Dup() + # A recoverable exception after Isend accepts a strict fan-out prefix + # cannot be repaired by another protocol message: an early recipient + # may already expose the decision. Capacity and packet construction are + # therefore checked before the first Isend, and all subsequent MPI + # failures terminate the communicator group instead of returning to + # Python with an ambiguous partial-send state. + try: + comm.Set_errhandler(MPI.ERRORS_ARE_FATAL) + rank = int(comm.Get_rank()) + except Exception: + # Dup() creates an owned communicator before any transport member + # can be published. Roll it back transactionally if the remaining + # MPI setup fails. + try: + comm.Free() + except Exception: + pass + raise + self._comm = comm + self.rank = rank + self._pending: deque[_PendingSend] = deque() + self._max_pending_sends = max_pending_sends + self._max_send_tests_per_progress = max_send_tests_per_progress + self._receive_sources = tuple(rank for rank in self.participants if rank != self.rank) + self._receive_cursor = 0 + self._closed = False + + def send(self, packet: _Packet, destination: int) -> None: + self.send_many(((packet, destination),)) + + def send_many(self, messages: Sequence[tuple[_Packet, int]]) -> None: + if self._closed: + raise RuntimeError("cannot send after consensus transport shutdown") + messages_tuple = tuple(messages) + for _, destination in messages_tuple: + if destination not in self.participants: + raise RuntimeError(f"consensus destination {destination} is not a participant") + self.progress() + projected_pending = len(self._pending) + len(messages_tuple) + if projected_pending > self._max_pending_sends: + kinds = sorted({packet.kind.name for packet, _ in messages_tuple}) + raise _ConsensusBackpressure( + "asynchronous consensus send backpressure limit exceeded: " + f"rank={self.rank}, batch={len(messages_tuple)}, kinds={kinds}, " + f"pending={len(self._pending)}, projected={projected_pending}, " + f"limit={self._max_pending_sends}" + ) + buffers = tuple(packet.encode() for packet, _ in messages_tuple) + for buffer, (_, destination) in zip(buffers, messages_tuple): + request = self._comm.Isend(buffer, dest=destination, tag=0) + self._pending.append(_PendingSend(buffer=buffer, request=request)) + + def progress(self) -> None: + # Bound every hot-path progress call. Incomplete requests rotate to the + # tail so a permanently slow send cannot starve later completions. + to_test = min(len(self._pending), self._max_send_tests_per_progress) + for _ in range(to_test): + pending = self._pending.popleft() + if not pending.request.Test(): + self._pending.append(pending) + + def receive(self, limit: int) -> list[_Packet]: + packets: list[_Packet] = [] + if limit <= 0: + return packets + if not self._receive_sources: + return packets + # Probe each source at most once per sweep, then rotate the first source + # for the next sweep. A continuously busy low rank therefore cannot + # starve a higher rank when ``limit`` is smaller than the world size. + while len(packets) < limit: + made_progress = False + source_count = len(self._receive_sources) + start = self._receive_cursor + for offset in range(source_count): + if len(packets) >= limit: + break + source_index = (start + offset) % source_count + source = self._receive_sources[source_index] + if not self._comm.Iprobe(source=source, tag=0): + continue + fields = np.empty(_PACKET_FIELDS, dtype=np.uint64) + self._comm.Recv(fields, source=source, tag=0) + packet = _Packet.decode(fields) + if packet.source != source: + raise RuntimeError( + f"consensus packet source mismatch: payload={packet.source}, mpi={source}" + ) + packets.append(packet) + made_progress = True + self._receive_cursor = (source_index + 1) % source_count + if not made_progress: + self._receive_cursor = (start + 1) % source_count + break + return packets + + @property + def pending_send_count(self) -> int: + return len(self._pending) + + def close(self, timeout_s: float) -> None: + if self._closed: + return + if timeout_s <= 0: + raise ValueError("transport close timeout must be positive") + deadline = time.monotonic() + timeout_s + while self._pending: + self.progress() + if self._pending and time.monotonic() >= deadline: + raise RuntimeError( + f"timed out draining {len(self._pending)} asynchronous consensus sends" + ) + if self._pending: + time.sleep(0.001) + self._comm.Free() + self._closed = True + + +_Key = tuple[ConsensusPhase, int, int] +_IntentKey = tuple[_MessageKind, _Key] +_UNKNOWN_FAILURE_KEY: _Key = (ConsensusPhase.TERMINAL, 0, 0) + + +@dataclass(frozen=True) +class _LocalIntent: + packet: _Packet + destination: int + + +_NORMAL_LOCAL_INTENT_KINDS = ( + _MessageKind.VOTE, + _MessageKind.WITHDRAW, + _MessageKind.READY_ACK, + _MessageKind.READY_ABORT_ACK, + _MessageKind.READY_ACTIVATE_ACK, +) +_COORDINATOR_MESSAGE_KINDS = { + _MessageKind.READY_PREPARE, + _MessageKind.READY_RELEASE, + _MessageKind.READY_COMPLETE, + _MessageKind.READY_ABORT, + _MessageKind.READY_ABORT_FINALIZE, + _MessageKind.TERMINAL_COMMIT, + _MessageKind.CLOSE_ACK, +} +_READY_OUTCOME_MESSAGE_KINDS = { + _MessageKind.READY_PREPARE, + _MessageKind.READY_ACK, + _MessageKind.READY_RELEASE, + _MessageKind.READY_ACTIVATE_ACK, + _MessageKind.READY_COMPLETE, +} +_WITHDRAWN_OUTCOME_MESSAGE_KINDS = { + _MessageKind.WITHDRAW, + _MessageKind.READY_ABORT, + _MessageKind.READY_ABORT_ACK, + _MessageKind.READY_ABORT_FINALIZE, + _MessageKind.CLOSE, + _MessageKind.CLOSE_ACK, +} +_PRIORITY_LOCAL_INTENT_LIMIT = 2 + + +class AsyncConsensusCoordinator: + """Asynchronous authoritative agreement over one PP participant domain.""" + + def __init__( + self, + transport: ConsensusTransport, + *, + scheduling_rank: int | None = None, + max_messages_per_poll: int = _DEFAULT_MAX_MESSAGES_PER_POLL, + max_completed_epochs: int = _DEFAULT_MAX_COMPLETED_EPOCHS, + max_open_rounds: int = _DEFAULT_MAX_OPEN_ROUNDS, + round_timeout_s: float = _DEFAULT_ROUND_TIMEOUT_S, + ready_lease_timeout_s: float | None = None, + clock: Callable[[], float] = time.monotonic, + ): + if len(transport.participants) < 2: + raise ValueError("asynchronous consensus requires at least two participants") + if transport.rank not in transport.participants: + raise ValueError("local rank is not in the consensus participant domain") + if max_messages_per_poll <= 0: + raise ValueError("max_messages_per_poll must be positive") + if max_completed_epochs <= 0: + raise ValueError("max_completed_epochs must be positive") + if max_open_rounds <= 0: + raise ValueError("max_open_rounds must be positive") + if round_timeout_s <= 0: + raise ValueError("round_timeout_s must be positive") + if ready_lease_timeout_s is not None and ready_lease_timeout_s <= 0: + raise ValueError("ready_lease_timeout_s must be positive when set") + self._transport = transport + self._participants = transport.participants + self._participant_set = set(self._participants) + self._coordinator_rank = self._participants[-1] + self._scheduling_rank = ( + self._participants[0] if scheduling_rank is None else int(scheduling_rank) + ) + if self._scheduling_rank not in self._participant_set: + raise ValueError("scheduling rank is not in the consensus participant domain") + self._max_messages_per_poll = max_messages_per_poll + self._max_completed_epochs = max_completed_epochs + self._max_open_rounds = max_open_rounds + self._round_timeout_s = round_timeout_s + # A fully prepared READY round can legitimately remain silent while + # rank zero waits for scheduler/KV capacity. That interval is an + # application lease, not missing protocol progress, so it must not + # inherit the short vote/ack watchdog. By default the executor's own + # request and hang policies bound it; deployments may opt into a + # distinct lease deadline without weakening the protocol watchdog. + self._ready_lease_timeout_s = ready_lease_timeout_s + self._clock = clock + + self._local_votes: dict[_Key, ConsensusOutcome] = {} + self._votes: dict[_Key, dict[int, ConsensusOutcome]] = {} + self._ready_required_acks: dict[_Key, set[int]] = {} + self._ready_acks: dict[_Key, set[int]] = {} + self._ready_activation_required_acks: dict[_Key, set[int]] = {} + self._ready_activation_acks: dict[_Key, set[int]] = {} + self._ready_abort_required_acks: dict[_Key, set[int]] = {} + self._ready_abort_acks: dict[_Key, set[int]] = {} + self._ready_abort_requested: set[_Key] = set() + self._local_ready_prepared: set[_Key] = set() + self._local_ready_acknowledged: set[_Key] = set() + self._local_ready_released: set[_Key] = set() + self._local_ready_activated: set[_Key] = set() + self._local_ready_aborting: set[_Key] = set() + self._local_ready_abort_acknowledged: set[_Key] = set() + self._completed_epoch: OrderedDict[tuple[ConsensusPhase, int], int] = OrderedDict() + self._events: deque[ConsensusEvent] = deque() + self._round_progress: dict[_Key, _RoundProgress] = {} + self._round_deadlines: OrderedDict[_Key, float] = OrderedDict() + self._ready_lease_deadlines: OrderedDict[_Key, float] = OrderedDict() + self._rounds_by_request: dict[tuple[ConsensusPhase, int], set[_Key]] = {} + self._open_ready_epoch: dict[int, _Key] = {} + # Normal local intents are bounded by open rounds: protocol ordering + # permits at most one unsent intent per round. FAIL_STOP and CLOSE use a + # separate reserved queue so overload cannot prevent fail-closed + # propagation or shutdown. + self._local_outbox: OrderedDict[_IntentKey, _LocalIntent] = OrderedDict() + self._priority_local_outbox: OrderedDict[_IntentKey, _LocalIntent] = OrderedDict() + self._coordinator_actions: deque[tuple[_CoordinatorAction, _Key]] = deque() + self._queued_coordinator_actions: set[tuple[_CoordinatorAction, _Key]] = set() + self._fatal_key: _Key | None = None + self._fatal_error: str | None = None + self._fail_stop_propagated = False + + self._shutdown_started = False + self._closed_peers: set[int] = set() + self._close_sent = False + self._close_acknowledged = False + self._close_ack_sent = False + + @property + def rank(self) -> int: + return self._transport.rank + + @property + def coordinator_rank(self) -> int: + return self._coordinator_rank + + @property + def scheduling_rank(self) -> int: + """Rank that may transition the request after READY_RELEASE.""" + return self._scheduling_rank + + def publish_ready(self, request_id: int, epoch: int = 0) -> None: + self._publish( + ConsensusPhase.READY, + request_id, + epoch, + ConsensusOutcome.READY, + ) + + def withdraw_ready(self, request_id: int, epoch: int = 0) -> bool: + """Withdraw a readiness vote before this rank grants its lease. + + ``True`` means that this epoch is, or is becoming, aborted. ``False`` + means this rank already acknowledged READY_PREPARE (or observed the + epoch's final decision), so the caller must let readiness finish and + treat cancellation as part of the request's next lifecycle phase. + """ + self._check_running() + key = (ConsensusPhase.READY, int(request_id), int(epoch)) + current = self._local_votes.get(key) + if current == ConsensusOutcome.WITHDRAWN: + return True + if key in self._local_ready_aborting: + return True + if key in self._local_ready_acknowledged: + return False + if self._is_stale(key): + return False + self._ensure_no_open_ready_epoch(key) + self._reserve_local_round(key) + packet = self._packet( + _MessageKind.WITHDRAW, + key, + ConsensusOutcome.WITHDRAWN, + ) + if self.rank == self._coordinator_rank: + self._local_votes[key] = ConsensusOutcome.WITHDRAWN + self._touch_round(key) + self._request_ready_abort(key) + else: + vote_intent_key = (_MessageKind.VOTE, key) + # A cancellation may arrive before an unsent readiness vote reaches + # MPI. The coordinator accepts withdrawal without a preceding local + # vote, so replace that intent rather than consuming two outbox + # credits or later resurrecting the round. + self._local_outbox.pop(vote_intent_key, None) + self._enqueue_local_intent(packet, self._coordinator_rank) + self._local_votes[key] = ConsensusOutcome.WITHDRAWN + self._touch_round(key) + self._drain_local_outbox(1) + return True + + def publish_terminal( + self, + request_id: int, + outcome: ConsensusOutcome, + epoch: int = 0, + ) -> None: + if outcome not in ( + ConsensusOutcome.COMPLETED, + ConsensusOutcome.FAILED, + ConsensusOutcome.CANCELLED, + ): + raise ValueError(f"invalid terminal consensus outcome: {outcome}") + self._publish(ConsensusPhase.TERMINAL, request_id, epoch, outcome) + + def acknowledge_ready(self, request_id: int, epoch: int = 0) -> None: + """Grant an irrevocable local lease after applying READY_PREPARE. + + Once this method succeeds, :meth:`withdraw_ready` cannot abort on + behalf of this rank. The request remains prepared until it is selected + by the authoritative PP schedule or an abort is finalized. + """ + self._check_running() + key = (ConsensusPhase.READY, int(request_id), int(epoch)) + if self._local_votes.get(key) != ConsensusOutcome.READY: + raise RuntimeError(f"cannot acknowledge unpublished readiness for {key}") + if key not in self._local_ready_prepared: + raise RuntimeError(f"cannot acknowledge readiness before READY_PREPARE for {key}") + if key in self._local_ready_aborting: + raise RuntimeError(f"cannot acknowledge readiness while aborting {key}") + if key in self._local_ready_acknowledged: + return + packet = self._packet(_MessageKind.READY_ACK, key, ConsensusOutcome.READY) + if self.rank == self._coordinator_rank: + self._local_ready_acknowledged.add(key) + self._touch_round(key) + self._record_ready_ack(key, self.rank) + else: + self._enqueue_local_intent(packet, self._coordinator_rank) + self._local_ready_acknowledged.add(key) + self._touch_round(key) + self._drain_local_outbox(1) + + def acknowledge_ready_abort(self, request_id: int, epoch: int = 0) -> None: + """Acknowledge that the local READY_ABORT rollback has been applied.""" + self._check_running() + key = (ConsensusPhase.READY, int(request_id), int(epoch)) + if key not in self._local_ready_aborting: + raise RuntimeError(f"cannot acknowledge readiness abort before READY_ABORT for {key}") + if key in self._local_ready_abort_acknowledged: + return + packet = self._packet( + _MessageKind.READY_ABORT_ACK, + key, + ConsensusOutcome.WITHDRAWN, + ) + if self.rank == self._coordinator_rank: + self._local_ready_abort_acknowledged.add(key) + self._touch_round(key) + self._record_ready_abort_ack(key, self.rank) + else: + self._enqueue_local_intent(packet, self._coordinator_rank) + self._local_ready_abort_acknowledged.add(key) + self._touch_round(key) + self._drain_local_outbox(1) + + def acknowledge_ready_activation(self, request_id: int, epoch: int = 0) -> None: + """Acknowledge activation by the authoritative PP schedule. + + PREPARE keeps follower requests hidden. Rank zero receives + READY_RELEASE after every prepare ACK and can select the request. The + resulting PP schedule is the activation token: each rank calls this + method only after the same scheduled request is locally visible. The + coordinator retains the epoch until every participant acknowledges. + """ + self._check_running() + key = (ConsensusPhase.READY, int(request_id), int(epoch)) + if key not in self._local_ready_acknowledged: + raise RuntimeError( + f"cannot activate readiness before PREPARE acknowledgement for {key}" + ) + if self.rank == self._scheduling_rank and key not in self._local_ready_released: + raise RuntimeError( + f"scheduling rank activated readiness before READY_RELEASE for {key}" + ) + if key in self._local_ready_activated: + return + packet = self._packet(_MessageKind.READY_ACTIVATE_ACK, key, ConsensusOutcome.READY) + if self.rank == self._coordinator_rank: + self._local_ready_activated.add(key) + self._touch_round(key) + self._record_ready_activation_ack(key, self.rank) + else: + self._enqueue_local_intent(packet, self._coordinator_rank) + self._local_ready_activated.add(key) + self._touch_round(key) + self._drain_local_outbox(1) + + def poll(self) -> list[ConsensusEvent]: + try: + self._transport.progress() + except RuntimeError as error: + first_failure = self._fail_stop_from_runtime_error( + _UNKNOWN_FAILURE_KEY, + error, + ) + if not first_failure: + self._raise_if_fatal() + raise + self._drain_local_outbox(self._max_messages_per_poll) + try: + packets = self._transport.receive(self._max_messages_per_poll) + except RuntimeError as error: + first_failure = self._fail_stop_from_runtime_error( + _UNKNOWN_FAILURE_KEY, + error, + ) + if not first_failure: + self._raise_if_fatal() + raise + for packet in packets: + try: + self._handle_packet(packet) + except RuntimeError as error: + first_failure = self._fail_stop_from_runtime_error( + (packet.phase, packet.request_id, packet.epoch), + error, + ) + if not first_failure: + self._raise_if_fatal() + raise + if self.rank == self._coordinator_rank: + self._advance_coordinator(self._max_messages_per_poll) + if not self._shutdown_started: + self._check_round_watchdogs() + self._raise_if_fatal() + events = list(self._events) + self._events.clear() + return events + + def shutdown(self, timeout_s: float = 30.0) -> None: + if timeout_s <= 0: + raise ValueError("shutdown timeout must be positive") + if self._close_acknowledged: + self._transport.close(timeout_s) + return + self._shutdown_started = True + self._local_outbox.clear() + close_key = (ConsensusPhase.READY, 0, 0) + if self.rank == self._coordinator_rank: + self._closed_peers.add(self.rank) + elif not self._close_sent: + self._enqueue_local_intent( + self._packet(_MessageKind.CLOSE, close_key, ConsensusOutcome.WITHDRAWN), + self._coordinator_rank, + priority=True, + ) + self._close_sent = True + self._drain_local_outbox(1) + deadline = time.monotonic() + timeout_s + while not self._close_acknowledged: + self.poll() + if time.monotonic() >= deadline: + raise RuntimeError( + "timed out waiting for asynchronous consensus shutdown acknowledgement" + ) + time.sleep(0.001) + self._transport.close(max(0.001, deadline - time.monotonic())) + + def _publish( + self, + phase: ConsensusPhase, + request_id: int, + epoch: int, + outcome: ConsensusOutcome, + ) -> None: + self._check_running() + key = (phase, int(request_id), int(epoch)) + if self._is_stale(key): + raise RuntimeError(f"cannot publish a stale consensus epoch: {key}") + if phase == ConsensusPhase.READY: + self._ensure_no_open_ready_epoch(key) + if key in self._local_ready_aborting: + raise RuntimeError(f"cannot publish readiness while aborting {key}") + previous = self._local_votes.get(key) + if previous is not None: + if previous != outcome: + raise RuntimeError( + f"local consensus vote changed for {key}: {previous.name} -> {outcome.name}" + ) + return + self._reserve_local_round(key) + packet = self._packet(_MessageKind.VOTE, key, outcome) + if self.rank == self._coordinator_rank: + self._local_votes[key] = outcome + self._touch_round(key) + self._record_vote(key, self.rank, outcome) + else: + self._enqueue_local_intent(packet, self._coordinator_rank) + self._local_votes[key] = outcome + self._touch_round(key) + self._drain_local_outbox(1) + + def _enqueue_local_intent( + self, + packet: _Packet, + destination: int, + *, + priority: bool = False, + ) -> None: + intent_key = (packet.kind, (packet.phase, packet.request_id, packet.epoch)) + intent = _LocalIntent(packet=packet, destination=destination) + outbox = self._priority_local_outbox if priority else self._local_outbox + previous = outbox.get(intent_key) + if previous is not None: + if previous != intent: + raise RuntimeError(f"local consensus intent changed for {intent_key}") + return + limit = _PRIORITY_LOCAL_INTENT_LIMIT if priority else self._max_open_rounds + if len(outbox) >= limit: + raise RuntimeError( + "asynchronous consensus local-intent limit exceeded: " + f"rank={self.rank}, priority={int(priority)}, kind={packet.kind.name}, " + f"request_id={packet.request_id}, epoch={packet.epoch}, " + f"queued={len(outbox)}, limit={limit}" + ) + outbox[intent_key] = intent + + def _drain_local_outbox(self, limit: int) -> None: + if limit <= 0: + return + sent = 0 + for outbox in (self._priority_local_outbox, self._local_outbox): + while outbox and sent < limit: + intent_key, intent = next(iter(outbox.items())) + try: + self._transport.send(intent.packet, intent.destination) + except _ConsensusBackpressure: + return + except RuntimeError as error: + key = intent_key[1] + already_fatal = self._fatal_key is not None + self._enter_local_fail_stop( + key, + "asynchronous consensus transport failed while sending a durable " + f"local intent: rank={self.rank}, kind={intent.packet.kind.name}, " + f"phase={key[0].name}, request_id={key[1]}, epoch={key[2]}, " + f"error={error}", + notify_coordinator=True, + ) + if already_fatal: + if self._shutdown_started: + return + if self._fatal_error is not None: + raise RuntimeError(self._fatal_error) from error + raise + outbox.pop(intent_key) + sent += 1 + + def _enter_local_fail_stop( + self, + key: _Key, + diagnostic: str, + *, + notify_coordinator: bool, + ) -> None: + if self._fatal_key is not None: + return + self._fatal_key = key + self._fatal_error = diagnostic + self._local_outbox.clear() + self._events.clear() + if notify_coordinator and self.rank != self._coordinator_rank: + self._enqueue_local_intent( + self._packet(_MessageKind.FAIL_STOP, key, ConsensusOutcome.FAILED), + self._coordinator_rank, + priority=True, + ) + + def _fail_stop_from_runtime_error(self, key: _Key, error: RuntimeError) -> bool: + """Record the first local protocol/transport error and notify peers. + + Returns ``True`` only when ``error`` established the local fatal state. + A best-effort notification is attempted immediately, while its durable + intent/action remains queued if transport progress is unavailable. + """ + first_failure = self._fatal_key is None + if first_failure: + if self.rank == self._coordinator_rank: + self._request_fail_stop( + key, + reported_by=self.rank, + diagnostic=str(error), + ) + else: + self._enter_local_fail_stop( + key, + str(error), + notify_coordinator=True, + ) + try: + if self.rank == self._coordinator_rank: + self._advance_coordinator(1) + else: + self._drain_local_outbox(1) + except RuntimeError: + # The first diagnostic remains authoritative. FAIL_STOP is a + # reserved durable obligation and a later poll/shutdown retries it. + pass + return first_failure + + def _handle_packet(self, packet: _Packet) -> None: + if packet.source not in self._participant_set: + raise RuntimeError(f"packet source {packet.source} is not a participant") + self._validate_packet_contract(packet) + key = (packet.phase, packet.request_id, packet.epoch) + if packet.kind == _MessageKind.CLOSE: + self._require_coordinator() + self._closed_peers.add(packet.source) + return + if packet.kind == _MessageKind.CLOSE_ACK: + if packet.source != self._coordinator_rank: + raise RuntimeError("close acknowledgement did not come from coordinator") + self._close_acknowledged = True + return + if packet.kind == _MessageKind.FAIL_STOP: + self._require_outcome(packet, ConsensusOutcome.FAILED) + if self.rank == self._coordinator_rank: + self._request_fail_stop(key, reported_by=packet.source) + elif packet.source == self._coordinator_rank: + self._enter_local_fail_stop( + key, + ( + "asynchronous consensus entered coordinated fail-stop: " + f"rank={self.rank}, coordinator={self._coordinator_rank}, " + f"phase={key[0].name}, request_id={key[1]}, epoch={key[2]}" + ), + notify_coordinator=False, + ) + else: + raise RuntimeError( + "consensus fail-stop message did not come from the coordinator: " + f"source={packet.source}, coordinator={self._coordinator_rank}" + ) + return + if self._fatal_key is not None: + return + if self._is_stale(key): + return + if packet.kind == _MessageKind.VOTE: + self._require_coordinator() + if key in self._ready_abort_requested: + return + self._record_vote(key, packet.source, packet.outcome) + elif packet.kind == _MessageKind.WITHDRAW: + self._require_coordinator() + if packet.phase != ConsensusPhase.READY: + raise RuntimeError("only READY consensus can be withdrawn") + # MPI preserves per-source ordering. If this source's ACK was + # already observed, its lease is irrevocable and a late or buggy + # withdrawal cannot overturn the round. + if packet.source in self._ready_acks.get(key, set()): + return + self._request_ready_abort(key) + elif packet.kind == _MessageKind.READY_PREPARE: + self._require_phase(packet, ConsensusPhase.READY) + self._require_outcome(packet, ConsensusOutcome.READY) + if key in self._local_ready_aborting: + return + local_vote = self._local_votes.get(key) + if local_vote == ConsensusOutcome.WITHDRAWN: + # PREPARE and a local cancellation can cross in flight. The + # already-published withdrawal makes the coordinator abort the + # round; never acknowledge or resurrect it here. + return + if local_vote != ConsensusOutcome.READY: + raise RuntimeError(f"received READY_PREPARE without a local READY vote: {key}") + if key in self._local_ready_prepared: + return + self._touch_round(key) + self._local_ready_prepared.add(key) + self._events.append( + ConsensusEvent( + ConsensusEventKind.READY_PREPARE, + packet.request_id, + packet.epoch, + ConsensusOutcome.READY, + ) + ) + elif packet.kind == _MessageKind.READY_ACK: + self._require_phase(packet, ConsensusPhase.READY) + self._require_coordinator() + self._record_ready_ack(key, packet.source) + elif packet.kind == _MessageKind.READY_ACTIVATE_ACK: + self._require_phase(packet, ConsensusPhase.READY) + self._require_coordinator() + self._record_ready_activation_ack(key, packet.source) + elif packet.kind == _MessageKind.READY_RELEASE: + self._require_phase(packet, ConsensusPhase.READY) + self._require_outcome(packet, ConsensusOutcome.READY) + if self.rank != self._scheduling_rank: + raise RuntimeError("non-scheduling rank received READY_RELEASE") + if key in self._local_ready_released: + return + self._touch_round(key) + self._local_ready_released.add(key) + self._events.append( + ConsensusEvent( + ConsensusEventKind.READY_RELEASE, + packet.request_id, + packet.epoch, + ConsensusOutcome.READY, + ) + ) + elif packet.kind == _MessageKind.READY_COMPLETE: + self._require_phase(packet, ConsensusPhase.READY) + self._require_outcome(packet, ConsensusOutcome.READY) + self._complete_local(key) + self._events.append( + ConsensusEvent( + ConsensusEventKind.READY_COMPLETE, + packet.request_id, + packet.epoch, + ConsensusOutcome.READY, + ) + ) + elif packet.kind == _MessageKind.READY_ABORT: + self._require_phase(packet, ConsensusPhase.READY) + self._require_outcome(packet, ConsensusOutcome.WITHDRAWN) + self._apply_ready_abort(key) + elif packet.kind == _MessageKind.READY_ABORT_ACK: + self._require_phase(packet, ConsensusPhase.READY) + self._require_outcome(packet, ConsensusOutcome.WITHDRAWN) + self._require_coordinator() + self._record_ready_abort_ack(key, packet.source) + elif packet.kind == _MessageKind.READY_ABORT_FINALIZE: + self._require_phase(packet, ConsensusPhase.READY) + self._require_outcome(packet, ConsensusOutcome.WITHDRAWN) + if key not in self._local_ready_aborting: + raise RuntimeError(f"received READY_ABORT_FINALIZE before READY_ABORT for {key}") + self._complete_local(key) + self._events.append( + ConsensusEvent( + ConsensusEventKind.READY_ABORT_FINALIZE, + packet.request_id, + packet.epoch, + ConsensusOutcome.WITHDRAWN, + ) + ) + elif packet.kind == _MessageKind.TERMINAL_COMMIT: + self._require_phase(packet, ConsensusPhase.TERMINAL) + if packet.outcome not in ( + ConsensusOutcome.COMPLETED, + ConsensusOutcome.FAILED, + ConsensusOutcome.CANCELLED, + ): + raise RuntimeError( + f"message TERMINAL_COMMIT has invalid outcome {packet.outcome.name}" + ) + if self._local_votes.get(key) not in ( + ConsensusOutcome.COMPLETED, + ConsensusOutcome.FAILED, + ConsensusOutcome.CANCELLED, + ): + raise RuntimeError(f"received terminal commit before local quiescence: {key}") + self._complete_local(key) + self._events.append( + ConsensusEvent( + ConsensusEventKind.TERMINAL_COMMIT, + packet.request_id, + packet.epoch, + packet.outcome, + ) + ) + else: + raise RuntimeError(f"unhandled consensus message kind {packet.kind}") + + def _advance_coordinator(self, limit: int) -> None: + for _ in range(limit): + if not self._coordinator_actions: + break + action, key = self._coordinator_actions.popleft() + action_key = (action, key) + if action_key not in self._queued_coordinator_actions: + continue + if action != _CoordinatorAction.FAIL_STOP and self._is_stale(key): + self._queued_coordinator_actions.discard(action_key) + continue + try: + if action == _CoordinatorAction.VOTES_COMPLETE: + votes = self._votes.get(key) + if votes is not None and len(votes) == len(self._participants): + if key[0] == ConsensusPhase.TERMINAL: + outcome = self._reduce_terminal(votes.values()) + self._broadcast_terminal_commit(key, outcome) + self._votes.pop(key, None) + elif key not in self._ready_required_acks: + self._prepare_ready(key) + elif action == _CoordinatorAction.READY_ACKS_COMPLETE: + required = self._ready_required_acks.get(key) + if ( + required is not None + and key not in self._ready_abort_requested + and self._ready_acks.get(key, set()) == required + ): + self._release_ready(key) + elif action == _CoordinatorAction.READY_ACTIVATE_ACKS_COMPLETE: + required = self._ready_activation_required_acks.get(key) + if ( + required is not None + and self._ready_activation_acks.get(key, set()) == required + ): + self._complete_ready(key) + elif action == _CoordinatorAction.READY_ABORT_ACKS_COMPLETE: + required = self._ready_abort_required_acks.get(key) + if required is not None and self._ready_abort_acks.get(key, set()) == required: + self._finalize_ready_abort(key) + elif action == _CoordinatorAction.READY_ABORT_START: + if key in self._ready_abort_requested: + self._start_ready_abort(key) + elif action == _CoordinatorAction.FAIL_STOP: + if key == self._fatal_key and not self._fail_stop_propagated: + self._propagate_fail_stop(key) + else: + raise RuntimeError(f"unhandled coordinator action {action}") + except _ConsensusBackpressure: + self._coordinator_actions.appendleft(action_key) + return + except RuntimeError as error: + if action == _CoordinatorAction.FAIL_STOP: + self._coordinator_actions.appendleft(action_key) + if self._shutdown_started: + return + if self._fatal_error is not None: + raise RuntimeError(self._fatal_error) from error + raise + # A decision fan-out can fail after the transport has accepted + # a strict prefix. Never retry that original action: a duplicate + # or second partial decision could let ranks diverge. Atomically + # replace it with the reserved global FAIL_STOP obligation. + self._queued_coordinator_actions.discard(action_key) + self._fail_stop_from_runtime_error(key, error) + raise + self._queued_coordinator_actions.discard(action_key) + + if ( + self._shutdown_started + and not self._close_ack_sent + and self._closed_peers == self._participant_set + ): + close_key = (ConsensusPhase.READY, 0, 0) + close_messages = tuple( + ( + self._packet( + _MessageKind.CLOSE_ACK, + close_key, + ConsensusOutcome.WITHDRAWN, + ), + rank, + ) + for rank in self._participants + if rank != self.rank + ) + try: + self._transport.send_many(close_messages) + except _ConsensusBackpressure: + return + self._votes.clear() + self._ready_required_acks.clear() + self._ready_acks.clear() + self._ready_activation_required_acks.clear() + self._ready_activation_acks.clear() + self._ready_abort_required_acks.clear() + self._ready_abort_acks.clear() + self._ready_abort_requested.clear() + self._local_votes.clear() + self._round_progress.clear() + self._round_deadlines.clear() + self._ready_lease_deadlines.clear() + self._rounds_by_request.clear() + self._open_ready_epoch.clear() + self._local_outbox.clear() + self._priority_local_outbox.clear() + self._coordinator_actions.clear() + self._queued_coordinator_actions.clear() + self._close_ack_sent = True + self._close_acknowledged = True + + def _prepare_ready(self, key: _Key) -> None: + required = set(self._participants) + messages = tuple( + ( + self._packet( + _MessageKind.READY_PREPARE, + key, + ConsensusOutcome.READY, + ), + rank, + ) + for rank in self._participants + if rank != self.rank + ) + self._transport.send_many(messages) + self._ready_required_acks[key] = required + self._ready_acks[key] = set() + self._touch_round(key) + for rank in self._participants: + if rank == self.rank: + self._local_ready_prepared.add(key) + self._events.append( + ConsensusEvent( + ConsensusEventKind.READY_PREPARE, + key[1], + key[2], + ConsensusOutcome.READY, + ) + ) + + def _record_vote( + self, + key: _Key, + source: int, + outcome: ConsensusOutcome, + ) -> None: + if source not in self._participant_set: + raise RuntimeError(f"vote source {source} is not a participant") + if key[0] == ConsensusPhase.READY and outcome != ConsensusOutcome.READY: + raise RuntimeError(f"invalid readiness vote {outcome.name}") + if key[0] == ConsensusPhase.TERMINAL and outcome not in ( + ConsensusOutcome.COMPLETED, + ConsensusOutcome.FAILED, + ConsensusOutcome.CANCELLED, + ): + raise RuntimeError(f"invalid terminal vote {outcome.name}") + self._reserve_coordinator_round(key) + votes = self._votes.setdefault(key, {}) + previous = votes.get(source) + if previous is not None and previous != outcome: + raise RuntimeError( + f"participant {source} changed vote for {key}: {previous.name} -> {outcome.name}" + ) + votes[source] = outcome + self._touch_round(key) + if len(votes) == len(self._participants): + self._enqueue_coordinator_action(_CoordinatorAction.VOTES_COMPLETE, key) + + def _record_ready_ack(self, key: _Key, source: int) -> None: + if key in self._ready_abort_requested: + return + required = self._ready_required_acks.get(key) + if required is None or source not in required: + raise RuntimeError(f"unexpected readiness acknowledgement from {source} for {key}") + self._ready_acks[key].add(source) + self._touch_round(key) + if self._ready_acks[key] == required: + self._enqueue_coordinator_action(_CoordinatorAction.READY_ACKS_COMPLETE, key) + + def _record_ready_activation_ack(self, key: _Key, source: int) -> None: + required = self._ready_activation_required_acks.get(key) + if required is None or source not in required: + raise RuntimeError( + f"unexpected readiness activation acknowledgement from {source} for {key}" + ) + self._ready_activation_acks[key].add(source) + self._touch_round(key) + if self._ready_activation_acks[key] == required: + self._enqueue_coordinator_action( + _CoordinatorAction.READY_ACTIVATE_ACKS_COMPLETE, + key, + ) + + def _record_ready_abort_ack(self, key: _Key, source: int) -> None: + required = self._ready_abort_required_acks.get(key) + if required is None or source not in required: + raise RuntimeError( + f"unexpected readiness abort acknowledgement from {source} for {key}" + ) + self._ready_abort_acks[key].add(source) + self._touch_round(key) + if self._ready_abort_acks[key] == required: + self._enqueue_coordinator_action( + _CoordinatorAction.READY_ABORT_ACKS_COMPLETE, + key, + ) + + def _enqueue_coordinator_action(self, action: _CoordinatorAction, key: _Key) -> None: + action_key = (action, key) + if action_key in self._queued_coordinator_actions: + return + self._queued_coordinator_actions.add(action_key) + self._coordinator_actions.append(action_key) + + def _request_fail_stop( + self, + key: _Key, + reported_by: int, + diagnostic: str | None = None, + ) -> None: + if self.rank != self._coordinator_rank: + raise RuntimeError("only the coordinator can propagate consensus fail-stop") + if self._fatal_key is not None: + return + self._fatal_key = key + self._fatal_error = diagnostic or ( + "asynchronous consensus entered coordinated fail-stop: " + f"rank={self.rank}, coordinator={self._coordinator_rank}, " + f"reported_by={reported_by}, phase={key[0].name}, " + f"request_id={key[1]}, epoch={key[2]}" + ) + self._local_outbox.clear() + self._events.clear() + self._coordinator_actions.clear() + self._queued_coordinator_actions.clear() + action_key = (_CoordinatorAction.FAIL_STOP, key) + self._queued_coordinator_actions.add(action_key) + self._coordinator_actions.appendleft(action_key) + + def _propagate_fail_stop(self, key: _Key) -> None: + messages = tuple( + ( + self._packet(_MessageKind.FAIL_STOP, key, ConsensusOutcome.FAILED), + rank, + ) + for rank in self._participants + if rank != self.rank + ) + self._transport.send_many(messages) + self._fail_stop_propagated = True + + def _release_ready(self, key: _Key) -> None: + if self.rank == self._scheduling_rank: + raise RuntimeError("ready coordinator must differ from the scheduling rank") + self._transport.send( + self._packet(_MessageKind.READY_RELEASE, key, ConsensusOutcome.READY), + self._scheduling_rank, + ) + self._ready_activation_required_acks[key] = set(self._participants) + self._ready_activation_acks[key] = set() + self._touch_round(key) + + def _complete_ready(self, key: _Key) -> None: + messages = tuple( + ( + self._packet(_MessageKind.READY_COMPLETE, key, ConsensusOutcome.READY), + rank, + ) + for rank in self._participants + if rank != self.rank + ) + self._transport.send_many(messages) + self._complete_local(key) + self._events.append( + ConsensusEvent( + ConsensusEventKind.READY_COMPLETE, + key[1], + key[2], + ConsensusOutcome.READY, + ) + ) + self._votes.pop(key, None) + self._ready_required_acks.pop(key, None) + self._ready_acks.pop(key, None) + self._ready_activation_required_acks.pop(key, None) + self._ready_activation_acks.pop(key, None) + + def _request_ready_abort(self, key: _Key) -> None: + if self.rank != self._coordinator_rank: + raise RuntimeError("only the coordinator can abort readiness") + if self._is_stale(key) or key in self._ready_abort_requested: + return + self._reserve_coordinator_round(key) + self._ready_abort_requested.add(key) + self._touch_round(key) + self._queued_coordinator_actions.discard((_CoordinatorAction.VOTES_COMPLETE, key)) + self._queued_coordinator_actions.discard((_CoordinatorAction.READY_ACKS_COMPLETE, key)) + self._enqueue_coordinator_action(_CoordinatorAction.READY_ABORT_START, key) + + def _start_ready_abort(self, key: _Key) -> None: + messages = tuple( + ( + self._packet( + _MessageKind.READY_ABORT, + key, + ConsensusOutcome.WITHDRAWN, + ), + rank, + ) + for rank in self._participants + if rank != self.rank + ) + self._transport.send_many(messages) + self._ready_abort_required_acks[key] = set(self._participants) + self._ready_abort_acks[key] = set() + self._touch_round(key) + self._votes.pop(key, None) + self._ready_required_acks.pop(key, None) + self._ready_acks.pop(key, None) + self._ready_activation_required_acks.pop(key, None) + self._ready_activation_acks.pop(key, None) + self._apply_ready_abort(key) + + def _apply_ready_abort(self, key: _Key) -> None: + if key in self._local_ready_aborting: + return + # READY_ABORT supersedes every normal outbound intent for the open + # readiness round. In particular, a follower can receive an abort from + # another rank's withdrawal while its own VOTE or READY_ACK is still + # waiting for transport capacity. Discarding those obsolete intents + # preserves the one-unsent-intent-per-round bound; the abort ACK below + # becomes the round's sole remaining outbound obligation. + for kind in ( + _MessageKind.VOTE, + _MessageKind.WITHDRAW, + _MessageKind.READY_ACK, + ): + self._local_outbox.pop((kind, key), None) + self._local_ready_aborting.add(key) + self._touch_round(key) + self._events.append( + ConsensusEvent( + ConsensusEventKind.READY_ABORT, + key[1], + key[2], + ConsensusOutcome.WITHDRAWN, + ) + ) + + def _finalize_ready_abort(self, key: _Key) -> None: + if self.rank != self._coordinator_rank: + raise RuntimeError("only the coordinator can finalize readiness abort") + messages = tuple( + ( + self._packet( + _MessageKind.READY_ABORT_FINALIZE, + key, + ConsensusOutcome.WITHDRAWN, + ), + rank, + ) + for rank in self._participants + if rank != self.rank + ) + self._transport.send_many(messages) + self._complete_local(key) + self._events.append( + ConsensusEvent( + ConsensusEventKind.READY_ABORT_FINALIZE, + key[1], + key[2], + ConsensusOutcome.WITHDRAWN, + ) + ) + + def _broadcast_terminal_commit( + self, + key: _Key, + outcome: ConsensusOutcome, + ) -> None: + messages = tuple( + ( + self._packet(_MessageKind.TERMINAL_COMMIT, key, outcome), + rank, + ) + for rank in self._participants + if rank != self.rank + ) + self._transport.send_many(messages) + self._complete_local(key) + self._events.append( + ConsensusEvent( + ConsensusEventKind.TERMINAL_COMMIT, + key[1], + key[2], + outcome, + ) + ) + + @staticmethod + def _reduce_terminal(outcomes: Iterable[ConsensusOutcome]) -> ConsensusOutcome: + values = set(outcomes) + if ConsensusOutcome.CANCELLED in values: + return ConsensusOutcome.CANCELLED + if ConsensusOutcome.FAILED in values: + return ConsensusOutcome.FAILED + if values == {ConsensusOutcome.COMPLETED}: + return ConsensusOutcome.COMPLETED + raise RuntimeError(f"invalid terminal vote set: {values}") + + def _complete_local(self, key: _Key) -> None: + epoch_key = (key[0], key[1]) + self._completed_epoch[epoch_key] = max(key[2], self._completed_epoch.get(epoch_key, -1)) + self._completed_epoch.move_to_end(epoch_key) + while len(self._completed_epoch) > self._max_completed_epochs: + self._completed_epoch.popitem(last=False) + # Completing epoch N tombstones every older epoch for this request and + # phase. Purge partial coordinator state too, otherwise an incomplete + # old round can remain resident forever after a newer round commits. + request_rounds = self._rounds_by_request.get(epoch_key, set()) + candidates = tuple(candidate for candidate in request_rounds if candidate[2] <= key[2]) + for candidate in candidates: + for state in ( + self._local_votes, + self._votes, + self._ready_required_acks, + self._ready_acks, + self._ready_activation_required_acks, + self._ready_activation_acks, + self._ready_abort_required_acks, + self._ready_abort_acks, + ): + state.pop(candidate, None) + self._local_ready_prepared.discard(candidate) + self._local_ready_acknowledged.discard(candidate) + self._local_ready_released.discard(candidate) + self._local_ready_activated.discard(candidate) + self._local_ready_aborting.discard(candidate) + self._local_ready_abort_acknowledged.discard(candidate) + self._ready_abort_requested.discard(candidate) + self._round_progress.pop(candidate, None) + self._round_deadlines.pop(candidate, None) + self._ready_lease_deadlines.pop(candidate, None) + for action in _CoordinatorAction: + self._queued_coordinator_actions.discard((action, candidate)) + for kind in _NORMAL_LOCAL_INTENT_KINDS: + self._local_outbox.pop((kind, candidate), None) + request_rounds.discard(candidate) + if not request_rounds: + self._rounds_by_request.pop(epoch_key, None) + if key[0] == ConsensusPhase.READY: + open_ready = self._open_ready_epoch.get(key[1]) + if open_ready is not None and open_ready[2] <= key[2]: + self._open_ready_epoch.pop(key[1], None) + + def _reserve_local_round(self, key: _Key) -> None: + self._reserve_round(key, "local publication") + + def _reserve_coordinator_round(self, key: _Key) -> None: + self._reserve_round(key, "coordinator receive") + + def _reserve_round(self, key: _Key, operation: str) -> None: + if key in self._round_progress: + return + if key[0] == ConsensusPhase.READY: + open_key = self._open_ready_epoch.get(key[1]) + if open_key is not None and open_key != key: + raise RuntimeError( + "cannot reuse a readiness request ID before its prior epoch finalizes: " + f"open={open_key}, new={key}" + ) + if len(self._round_progress) >= self._max_open_rounds: + oldest_key = next(iter(self._round_progress), None) + raise RuntimeError( + "asynchronous consensus open-round limit exceeded: " + f"rank={self.rank}, operation={operation}, key={key}, " + f"open={len(self._round_progress)}, limit={self._max_open_rounds}, " + f"oldest={oldest_key}, pending_sends={self._transport.pending_send_count}" + ) + + def _touch_round(self, key: _Key) -> None: + now = self._clock() + progress = self._round_progress.get(key) + if progress is None: + self._reserve_round(key, "protocol progress") + self._round_progress[key] = _RoundProgress( + started_at=now, + last_progress_at=now, + ) + self._arm_round_watchdog(key, now) + self._rounds_by_request.setdefault((key[0], key[1]), set()).add(key) + if key[0] == ConsensusPhase.READY: + self._open_ready_epoch[key[1]] = key + return + progress.last_progress_at = now + # The watchdog measures lack of protocol progress, not total round + # lifetime. Readiness may legitimately remain open while the scheduler + # waits for capacity, and later votes/acknowledgements must grant a new + # idle window. Moving the key preserves deadline order for the O(1) + # oldest-round check below. + self._arm_round_watchdog(key, now) + + def _is_silent_ready_lease(self, key: _Key) -> bool: + if key[0] != ConsensusPhase.READY: + return False + if key in self._ready_abort_requested or key in self._local_ready_aborting: + return False + if self.rank == self._coordinator_rank: + required = self._ready_activation_required_acks.get(key) + return required is not None and not self._ready_activation_acks.get(key) + return key in self._local_ready_acknowledged and key not in self._local_ready_activated + + def _arm_round_watchdog(self, key: _Key, now: float) -> None: + self._round_deadlines.pop(key, None) + self._ready_lease_deadlines.pop(key, None) + if self._is_silent_ready_lease(key): + if self._ready_lease_timeout_s is not None: + self._ready_lease_deadlines[key] = now + self._ready_lease_timeout_s + return + self._round_deadlines[key] = now + self._round_timeout_s + + def _check_round_watchdogs(self) -> None: + candidates: list[tuple[_Key, float]] = [] + if self._round_deadlines: + candidates.append(next(iter(self._round_deadlines.items()))) + if self._ready_lease_deadlines: + candidates.append(next(iter(self._ready_lease_deadlines.items()))) + if not candidates: + return + key, deadline = min(candidates, key=lambda candidate: candidate[1]) + now = self._clock() + if now < deadline: + return + progress = self._round_progress[key] + votes = self._votes.get(key, {}) + required_acks = self._ready_required_acks.get(key, set()) + ready_acks = self._ready_acks.get(key, set()) + required_activation_acks = self._ready_activation_required_acks.get(key, set()) + activation_acks = self._ready_activation_acks.get(key, set()) + required_abort_acks = self._ready_abort_required_acks.get(key, set()) + abort_acks = self._ready_abort_acks.get(key, set()) + missing_votes = sorted(self._participant_set - set(votes)) + missing_ready_acks = sorted(required_acks - ready_acks) + missing_activation_acks = sorted(required_activation_acks - activation_acks) + missing_abort_acks = sorted(required_abort_acks - abort_acks) + local_vote = self._local_votes.get(key) + diagnostic = ( + "asynchronous consensus round watchdog expired without a global decision: " + f"rank={self.rank}, coordinator={self._coordinator_rank}, " + f"phase={key[0].name}, request_id={key[1]}, epoch={key[2]}, " + f"age_s={now - progress.started_at:.3f}, " + f"idle_s={now - progress.last_progress_at:.3f}, " + f"local_vote={local_vote.name if local_vote is not None else 'NONE'}, " + f"missing_votes={missing_votes}, missing_ready_acks={missing_ready_acks}, " + f"missing_activation_acks={missing_activation_acks}, " + f"missing_abort_acks={missing_abort_acks}, " + f"pending_sends={self._transport.pending_send_count}" + ) + if self.rank == self._coordinator_rank: + self._request_fail_stop( + key, + reported_by=self.rank, + diagnostic=diagnostic, + ) + self._advance_coordinator(1) + return + self._enter_local_fail_stop( + key, + diagnostic, + notify_coordinator=True, + ) + # Best effort only: local fail-stop must not depend on notification + # capacity. The reserved intent remains queued for shutdown progress if + # the transport is currently full. + self._drain_local_outbox(1) + + def _raise_if_fatal(self) -> None: + if self._fatal_error is not None and not self._shutdown_started: + raise RuntimeError(self._fatal_error) + + def _is_stale(self, key: _Key) -> bool: + return key[2] <= self._completed_epoch.get((key[0], key[1]), -1) + + def _ensure_no_open_ready_epoch(self, key: _Key) -> None: + open_key = self._open_ready_epoch.get(key[1]) + if open_key is not None and open_key != key: + raise RuntimeError( + "cannot reuse a readiness request ID before its prior epoch finalizes: " + f"open={open_key}, new={key}" + ) + + def _packet( + self, + kind: _MessageKind, + key: _Key, + outcome: ConsensusOutcome, + ) -> _Packet: + return _Packet(kind, key[0], key[1], key[2], outcome, self.rank) + + def _validate_packet_contract(self, packet: _Packet) -> None: + if packet.kind in _COORDINATOR_MESSAGE_KINDS and packet.source != self._coordinator_rank: + raise RuntimeError( + f"message {packet.kind.name} did not come from coordinator " + f"{self._coordinator_rank}: source={packet.source}" + ) + + if packet.kind in _READY_OUTCOME_MESSAGE_KINDS: + self._require_phase(packet, ConsensusPhase.READY) + self._require_outcome(packet, ConsensusOutcome.READY) + elif packet.kind in _WITHDRAWN_OUTCOME_MESSAGE_KINDS: + self._require_phase(packet, ConsensusPhase.READY) + self._require_outcome(packet, ConsensusOutcome.WITHDRAWN) + elif packet.kind == _MessageKind.TERMINAL_COMMIT: + self._require_phase(packet, ConsensusPhase.TERMINAL) + if packet.outcome not in ( + ConsensusOutcome.COMPLETED, + ConsensusOutcome.FAILED, + ConsensusOutcome.CANCELLED, + ): + raise RuntimeError( + f"message TERMINAL_COMMIT has invalid outcome {packet.outcome.name}" + ) + elif packet.kind == _MessageKind.FAIL_STOP: + self._require_outcome(packet, ConsensusOutcome.FAILED) + elif packet.kind == _MessageKind.VOTE: + if packet.phase == ConsensusPhase.READY: + self._require_outcome(packet, ConsensusOutcome.READY) + elif packet.outcome not in ( + ConsensusOutcome.COMPLETED, + ConsensusOutcome.FAILED, + ConsensusOutcome.CANCELLED, + ): + raise RuntimeError( + f"message VOTE has invalid terminal outcome {packet.outcome.name}" + ) + + def _require_coordinator(self) -> None: + if self.rank != self._coordinator_rank: + raise RuntimeError("only the coordinator may receive this consensus message") + + @staticmethod + def _require_phase(packet: _Packet, phase: ConsensusPhase) -> None: + if packet.phase != phase: + raise RuntimeError( + f"message {packet.kind.name} has phase {packet.phase.name}, expected {phase.name}" + ) + + @staticmethod + def _require_outcome(packet: _Packet, outcome: ConsensusOutcome) -> None: + if packet.outcome != outcome: + raise RuntimeError( + f"message {packet.kind.name} has outcome {packet.outcome.name}, " + f"expected {outcome.name}" + ) + + def _check_running(self) -> None: + if self._shutdown_started: + raise RuntimeError("cannot publish consensus state after shutdown starts") + if self._fatal_key is not None: + raise RuntimeError("cannot publish consensus state after coordinated fail-stop starts") diff --git a/tensorrt_llm/_torch/disaggregation/base/transfer.py b/tensorrt_llm/_torch/disaggregation/base/transfer.py index 320f7e067bdf..6fba7e02e484 100644 --- a/tensorrt_llm/_torch/disaggregation/base/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/base/transfer.py @@ -1,3 +1,18 @@ +# 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. + from __future__ import annotations from abc import ABC, abstractmethod @@ -144,6 +159,14 @@ def is_completed(self) -> bool: ... @abstractmethod def wait_complete(self, blocking: bool = False) -> Optional[WaitResult]: ... + @abstractmethod + def seal_and_check_quiescent(self) -> bool: + """Prevent new transfer work and report whether in-flight work has drained.""" + + @abstractmethod + def has_transferring_tasks(self) -> bool: + """Return whether the session has transfer work in flight.""" + @property @abstractmethod def exception(self) -> Optional[Exception]: ... diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/core.py b/tensorrt_llm/_torch/disaggregation/native/bounce/core.py index 1878fa90fd3a..52ebbca2d2a6 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/core.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/core.py @@ -29,7 +29,7 @@ class TransferState(Enum): SCATTERING = "scattering" # all writers succeeded; scattering back into the cache COMPLETED = "completed" # scattered cleanly, slot released FAILED = "failed" # all writers done, at least one failed, slot released - QUARANTINED = "quarantined" # a writer is in doubt, slot held out of reuse + CANCELLED_DRAINED = "cancelled_drained" # cancellation settled only after exact drain proof class ScatterState(Enum): @@ -60,22 +60,27 @@ class Settlement: @dataclass class TransferContext: """Lifetime state machine for one bounced receive region. The region is released only after every - writer reports (success or failure both mean its write has drained); a writer given up on (a - cancelled or timed-out transfer, whose one-sided write cannot be aborted) makes it quarantined - instead of reused.""" + writer reports (success or failure both mean its write has drained), or an exact sender ACK proves + that every advertised write for the request incarnation has drained. Cancellation alone never + makes the region reusable.""" rid_slice: Tuple[int, int] slot_id: int base_addr: int per_writer_bytes: int num_writers: int + allowed_destination_ranges: Tuple[Tuple[int, int], ...] = () + expected_destination_plans: Dict[int, Tuple[Tuple[int, int], ...]] = field(default_factory=dict) on_done: Optional[Callable[[bool], None]] = None + on_settled: Optional[Callable[[bool], None]] = None # Whether each writer that reported succeeded, keyed by rank; presence means it reported. _writer_ok: Dict[int, bool] = field(default_factory=dict) + _expected_writer_bases: Dict[int, int] = field(default_factory=dict) # per successful writer: where it wrote, plus the fragments to scatter back _scatter_descs: List[tuple] = field(default_factory=list) _orphaned: bool = False + _drain_proven: bool = False scatter_state: ScatterState = ScatterState.IDLE state: TransferState = TransferState.INIT settled: bool = False @@ -89,7 +94,7 @@ def _writers_final(self) -> bool: TransferState.SCATTERING, TransferState.COMPLETED, TransferState.FAILED, - TransferState.QUARANTINED, + TransferState.CANCELLED_DRAINED, ) def _all_writers_reported(self) -> bool: @@ -99,6 +104,40 @@ def _all_writers_succeeded(self) -> bool: return self._all_writers_reported() and all(self._writer_ok.values()) # Mutations: call only while holding the transport's reservation lock. + def bind_writer(self, peer_rank: int, src_base: int) -> None: + """Bind one immutable writer identity to its advertised bounce sub-region.""" + if self._writers_final() or peer_rank in self._expected_writer_bases: + raise RuntimeError(f"bounce writer {peer_rank} was bound more than once") + if len(self._expected_writer_bases) >= self.num_writers: + raise RuntimeError("bounce writer binding exceeds the reserved fan-in") + if self.expected_destination_plans and peer_rank not in self.expected_destination_plans: + raise RuntimeError(f"bounce writer {peer_rank} has no receiver-derived scatter plan") + self._expected_writer_bases[peer_rank] = src_base + + def set_completion_callback(self, on_settled: Callable[[bool], None]) -> None: + """Install the immutable reservation-lifetime callback before advertisement.""" + if self._writers_final() or self.on_settled is not None: + raise RuntimeError("bounce completion callback was installed more than once") + self.on_settled = on_settled + + def _combined_callback(self) -> Optional[Callable[[bool], None]]: + if self.on_done is None: + return self.on_settled + if self.on_settled is None: + return self.on_done + on_done = self.on_done + on_settled = self.on_settled + + def complete(success: bool) -> None: + try: + on_done(success) + finally: + # Session lifetime credit must retire even if optional task/perf + # completion handling raises. + on_settled(success) + + return complete + def record_writer_result( self, peer_rank: int, @@ -110,6 +149,24 @@ def record_writer_result( ) -> None: """Record one writer's terminal report. Repeat or late reports are ignored, so duplicate or out-of-order messages are harmless.""" + expected_src_base = self._expected_writer_bases.get(peer_rank) + if expected_src_base is None: + raise RuntimeError(f"result from unexpected bounce writer rank {peer_rank}") + if succeeded: + if src_base is None or dst_ptrs is None or sizes is None: + raise RuntimeError( + "successful bounce writer result requires a complete scatter tail" + ) + if src_base != expected_src_base: + raise RuntimeError( + f"bounce writer {peer_rank} returned source base {src_base}, " + f"expected {expected_src_base}" + ) + elif src_base is not None and src_base != expected_src_base: + raise RuntimeError( + f"bounce writer {peer_rank} returned source base {src_base}, " + f"expected {expected_src_base}" + ) # A duplicate or late report can still arrive after the writer set is final, because # scatter runs on another thread while the region stays live. For example a retransmitted # notification, or a stray failure that would flip a good transfer to failed; drop it. @@ -125,11 +182,16 @@ def record_writer_result( def mark_orphaned(self) -> None: """Give up on writers that never reported (cancel, timeout, shutdown): a one-sided write - cannot be aborted, so the region is quarantined on settle. No-op once scattering or done.""" + cannot be aborted, so retain the region pending explicit drain proof. No-op once scattering + or done.""" if self._writers_final(): return self._orphaned = True + def confirm_drained(self) -> None: + """Record external proof that no writer can touch this region again.""" + self._drain_proven = True + def begin_scatter(self) -> None: self.state = TransferState.SCATTERING self.scatter_state = ScatterState.QUEUED @@ -154,7 +216,10 @@ def ready_to_settle(self) -> bool: if self.settled: return False if self._orphaned: - return True # in doubt: settle now and quarantine + # Cancellation alone says nothing about a one-sided remote write. + # Settle only after every writer reports terminal or the sender's + # drain ACK provides equivalent proof. + return self._drain_proven or self._all_writers_reported() if not self._all_writers_reported(): return False # a writer has not reported yet if self._all_writers_succeeded() and self._scatter_descs: @@ -168,11 +233,23 @@ def settle(self) -> Optional[Settlement]: return None self.settled = True if self._orphaned: - self.state = TransferState.QUARANTINED - return Settlement(self.slot_id, Disposition.QUARANTINE, False, self.on_done) + self.state = TransferState.CANCELLED_DRAINED + # Despite the historical state name, reaching here now requires + # exact drain proof, so the region is safe to release immediately. + return Settlement( + self.slot_id, + Disposition.RELEASE, + False, + self._combined_callback(), + ) success = self._all_writers_succeeded() and self.scatter_state is not ScatterState.FAILED self.state = TransferState.COMPLETED if success else TransferState.FAILED - return Settlement(self.slot_id, Disposition.RELEASE, success, self.on_done) + return Settlement( + self.slot_id, + Disposition.RELEASE, + success, + self._combined_callback(), + ) class BounceTransport(ABC): @@ -192,13 +269,28 @@ def release_send(self, slot_id) -> None: """Release a send region after its write completes.""" @abstractmethod - def reserve(self, recv_req, num_writers: int = 1, *, timeout: Optional[float] = None) -> bool: + def reserve( + self, + recv_req, + num_writers: int = 1, + *, + timeout: Optional[float] = None, + expected_destination_plans=None, + ) -> bool: """Reserve a region and record its address for the senders. False falls back to per-fragment.""" @abstractmethod def writer_base(self, rid_slice, writer_index: int) -> Optional[int]: """Where the given fan-in writer writes in the region.""" + @abstractmethod + def bind_writer(self, rid_slice, peer_rank: int, writer_index: int) -> Optional[int]: + """Bind a rank to its immutable advertised sub-region and return that base.""" + + @abstractmethod + def set_completion_callback(self, rid_slice, on_settled: Callable[[bool], None]) -> None: + """Install the callback that retires the reservation's lifetime credit.""" + @abstractmethod def is_bounced(self, rid_slice) -> bool: """Whether this request and slice name a live bounced region.""" @@ -209,7 +301,11 @@ def release_idle_reservation(self, rid_slice) -> None: @abstractmethod def orphan_reservation(self, rid_slice) -> None: - """Give up on an in-flight reservation (cancel/timeout/lost result); quarantine, don't leak.""" + """Retain an in-flight reservation until explicit drain proof.""" + + @abstractmethod + def confirm_drained(self, rid_slice) -> None: + """Release an orphan only after all possible remote writes drained.""" @abstractmethod def record_result( @@ -218,7 +314,7 @@ def record_result( """Handle a writer's success; scatter and finalize once all writers reported.""" @abstractmethod - def record_failure(self, rid_slice, peer_rank) -> None: + def record_failure(self, rid_slice, peer_rank, on_done=None) -> None: """Handle a writer's failure; free the region once all writers reported.""" @abstractmethod diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py index a9597b7cbb2a..c9745a3fadf8 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py @@ -57,7 +57,14 @@ class VmmBounceTransport(BounceTransport): @classmethod def from_config( - cls, agent, cfg, *, device_id: int, block_bytes_per_group: List[int] + cls, + agent, + cfg, + *, + device_id: int, + block_bytes_per_group: List[int], + destination_pool_layouts: Optional[List[List[tuple[int, int, int]]]] = None, + valid_destination_ranges: Optional[List[tuple[int, int]]] = None, ) -> Optional["VmmBounceTransport"]: """Build a transport sized from the config and clamped to free memory, or None if not even one chunk fits.""" @@ -83,6 +90,8 @@ def from_config( capacity_bytes=capacity_bytes, phys_chunk_size=chunk, block_bytes_per_group=block_bytes_per_group, + destination_pool_layouts=destination_pool_layouts, + valid_destination_ranges=valid_destination_ranges, min_blocks=cfg.min_blocks, ) @@ -94,6 +103,8 @@ def __init__( capacity_bytes: int, phys_chunk_size: int, block_bytes_per_group: List[int], + destination_pool_layouts: Optional[List[List[tuple[int, int, int]]]] = None, + valid_destination_ranges: Optional[List[tuple[int, int]]] = None, min_blocks: int = 96, quarantine_grace_s: float = _QUARANTINE_GRACE_S, name: str = "kv_bounce", @@ -102,6 +113,16 @@ def __init__( self._device_id = device_id # The byte size of one cache block, listed for each attention layer group. self._block_bytes_per_group = list(block_bytes_per_group) + # Receiver-owned physical slot layouts, indexed by layer group. These + # produce a request-specific destination allowlist and validate the + # exact per-rank plans before any bounce address is advertised. + # ``valid_destination_ranges`` is retained only as a compatibility/testing + # fallback for callers without page-table metadata; production + # construction always supplies the layouts. + self._destination_pool_layouts = tuple( + tuple(group) for group in (destination_pool_layouts or ()) + ) + self._valid_destination_ranges = tuple(valid_destination_ranges or ()) # Below this many blocks, skip bounce: coalescing only pays off for long context (the default # is roughly twelve thousand tokens; a heuristic, and tunable). self._min_blocks = min_blocks @@ -195,14 +216,45 @@ def _reserve_and_gather(self, write_meta, *, timeout): raise return slot_id, src_addr, total + @staticmethod + def _canonicalize_write_fragments(write_meta) -> None: + """Put bounce fragments in receiver-verifiable destination order. + + The contiguous bounce source preserves this order. Sorting all three + arrays together therefore keeps source data paired with its destination + while making the result tail independently checkable by the receiver. + """ + if write_meta.dst_ptrs.size < 2: + return + order = sorted( + range(write_meta.dst_ptrs.size), + key=lambda index: int(write_meta.dst_ptrs[index]), + ) + if order == list(range(write_meta.dst_ptrs.size)): + return + write_meta.src_ptrs = np.asarray( + [write_meta.src_ptrs[index] for index in order], dtype=np.int64 + ) + write_meta.dst_ptrs = np.asarray( + [write_meta.dst_ptrs[index] for index in order], dtype=np.int64 + ) + write_meta.sizes = np.asarray([write_meta.sizes[index] for index in order], dtype=np.int64) + def build_request(self, write_meta): """Gather into a send slot and build the coalesced write, or None on backpressure. The gather blocks (and frees the slot on failure) inside _reserve_and_gather.""" + self._canonicalize_write_fragments(write_meta) gathered = self._reserve_and_gather(write_meta, timeout=_RESERVE_TIMEOUT_S) if gathered is None: # backpressure: fall back return None slot_id, src_addr, total = gathered - return self._make_write(src_addr, write_meta, total), slot_id + try: + return self._make_write(src_addr, write_meta, total), slot_id + except Exception: + # The gather completed but submission has not started, so this + # send slot is still locally quiescent and can be released. + self._send_alloc.release(slot_id) + raise def release_send(self, slot_id) -> None: """Release a send region after its write has completed.""" @@ -217,16 +269,41 @@ def _skip_bounce(reason: str, *, warn_key: Optional[str] = None) -> bool: return False def reserve( - self, recv_req, num_writers: int = 1, *, timeout: Optional[float] = _RESERVE_TIMEOUT_S + self, + recv_req, + num_writers: int = 1, + *, + timeout: Optional[float] = _RESERVE_TIMEOUT_S, + expected_destination_plans: Optional[dict[int, tuple[np.ndarray, np.ndarray]]] = None, ) -> bool: """Reserve a region and create its state, recording the address for the senders. Returns False to fall back to the per-fragment path. A fan-in splits the region evenly, so the total must divide across the writers.""" - nblocks = sum(int(a.size) for a in recv_req.block_ids_per_layer_groups) + if getattr(recv_req, "mamba_state_index", None) is not None: + return self._skip_bounce("mamba state has no receiver-owned bounce scatter plan") + + valid_block_ids_per_group: list[np.ndarray] = [] + for block_ids in recv_req.block_ids_per_layer_groups: + # Production requests carry int64 ndarrays. The size-only fallback + # keeps lightweight unit fakes usable while production remains + # request-bound through ``_destination_pool_layouts`` below. + try: + values = np.asarray(block_ids, dtype=np.int64) + except (TypeError, ValueError): + values = np.asarray(list(range(int(block_ids.size))), dtype=np.int64) + if values.ndim != 1: + return self._skip_bounce("receiver block IDs are not one-dimensional") + values = np.asarray( + [int(block_id) for block_id in values if int(block_id) >= 0], + dtype=np.int64, + ) + valid_block_ids_per_group.append(values) + + nblocks = sum(int(a.size) for a in valid_block_ids_per_group) if nblocks < self._min_blocks: return self._skip_bounce(f"{nblocks} blocks < min {self._min_blocks} (too small)") total = 0 - for g, block_ids in enumerate(recv_req.block_ids_per_layer_groups): + for g, block_ids in enumerate(valid_block_ids_per_group): if g >= len(self._block_bytes_per_group): return self._skip_bounce(f"layer group {g} has no known slot size (e.g. mamba)") total += int(block_ids.size) * self._block_bytes_per_group[g] @@ -244,7 +321,7 @@ def reserve( # when the per-block sizes match, so require that here, else fall back. present_slot_bytes = { self._block_bytes_per_group[g] - for g, block_ids in enumerate(recv_req.block_ids_per_layer_groups) + for g, block_ids in enumerate(valid_block_ids_per_group) if int(block_ids.size) > 0 } if len(present_slot_bytes) > 1: @@ -266,6 +343,18 @@ def reserve( ) slot_id, addr = res recv_req.bounce_dst_base = addr + try: + allowed_destination_ranges = self._request_destination_ranges(valid_block_ids_per_group) + normalized_destination_plans = self._normalize_destination_plans( + expected_destination_plans, + num_writers=num_writers, + per_writer_bytes=total // num_writers, + allowed_destination_ranges=allowed_destination_ranges, + ) + except ValueError as error: + self._recv_alloc.release(slot_id) + recv_req.bounce_dst_base = None + return self._skip_bounce(str(error)) with self._reserved_map_lock: ctx = TransferContext( rid_slice=(recv_req.unique_rid, recv_req.slice_id), @@ -273,6 +362,8 @@ def reserve( base_addr=addr, per_writer_bytes=total // num_writers, num_writers=num_writers, + allowed_destination_ranges=allowed_destination_ranges, + expected_destination_plans=normalized_destination_plans, ) self._reserved_map[ctx.rid_slice] = ctx # inactive until the first writer reports # Positive marker: all fall-back guards above passed, so this transfer provably takes the @@ -285,12 +376,137 @@ def reserve( ) return True + def _request_destination_ranges( + self, block_ids_per_group: List[np.ndarray] + ) -> tuple[tuple[int, int], ...]: + """Build the immutable receiver-owned destination allowlist. + + Result tails describe scatter fragments for performance, but they must + never gain authority to write into another request's cache blocks. The + receiver derives the only legal physical slots from its own page table + and the block IDs already attached to this receive task. + """ + if not self._destination_pool_layouts: + if not self._valid_destination_ranges: + raise ValueError("no receiver-owned KV destination ranges are available") + return self._valid_destination_ranges + if len(block_ids_per_group) > len(self._destination_pool_layouts): + raise ValueError("receiver block groups exceed the local KV pool layout") + + ranges: set[tuple[int, int]] = set() + for group_index, block_ids in enumerate(block_ids_per_group): + layouts = self._destination_pool_layouts[group_index] + if block_ids.size and not layouts: + raise ValueError(f"layer group {group_index} has no local KV destination pool") + for base, slot_bytes, num_slots in layouts: + bad = next( + (int(block_id) for block_id in block_ids if int(block_id) >= num_slots), + None, + ) + if bad is not None: + raise ValueError( + f"receiver block ID {bad} exceeds layer-group {group_index} " + f"pool capacity {num_slots}" + ) + ranges.update( + (base + int(block_id) * slot_bytes, base + (int(block_id) + 1) * slot_bytes) + for block_id in block_ids + ) + if not ranges: + raise ValueError("receiver request has no valid KV destination slots") + return tuple(sorted(ranges)) + + def _normalize_destination_plans( + self, + plans: Optional[dict[int, tuple[np.ndarray, np.ndarray]]], + *, + num_writers: int, + per_writer_bytes: int, + allowed_destination_ranges: tuple[tuple[int, int], ...], + ) -> dict[int, tuple[tuple[int, int], ...]]: + """Validate and freeze exact receiver-derived plans before advertisement.""" + if not self._destination_pool_layouts: + return {} + if plans is None or len(plans) != num_writers: + raise ValueError( + "receiver-derived bounce destination plans do not match the writer fan-in" + ) + + normalized: dict[int, tuple[tuple[int, int], ...]] = {} + all_fragments: list[tuple[int, int]] = [] + for peer_rank, (dst_ptrs, sizes) in plans.items(): + dst_ptrs = np.asarray(dst_ptrs, dtype=np.int64) + sizes = np.asarray(sizes, dtype=np.int64) + if dst_ptrs.ndim != 1 or sizes.ndim != 1 or dst_ptrs.size != sizes.size: + raise ValueError(f"invalid receiver-derived scatter plan for rank {peer_rank}") + if np.any(sizes <= 0): + raise ValueError( + f"receiver-derived scatter plan for rank {peer_rank} has non-positive sizes" + ) + plan = tuple((int(ptr), int(size)) for ptr, size in zip(dst_ptrs, sizes, strict=True)) + if sum(size for _ptr, size in plan) != per_writer_bytes: + raise ValueError( + f"receiver-derived scatter plan for rank {peer_rank} does not describe " + f"exactly {per_writer_bytes} bytes" + ) + for index, (ptr, size) in enumerate(plan): + end = ptr + size + if index and ptr < plan[index - 1][0]: + raise ValueError( + f"receiver-derived scatter plan for rank {peer_rank} is not canonical" + ) + if index and ptr < plan[index - 1][0] + plan[index - 1][1]: + raise ValueError(f"receiver-derived scatter plan for rank {peer_rank} overlaps") + if not any( + valid_start <= ptr and end <= valid_end + for valid_start, valid_end in allowed_destination_ranges + ): + raise ValueError( + f"receiver-derived scatter plan for rank {peer_rank} is outside " + "the request's KV slots" + ) + all_fragments.append((ptr, end)) + normalized[int(peer_rank)] = plan + + try: + actual = self._coalesce_destination_ranges(all_fragments, reject_overlap=True) + expected = self._coalesce_destination_ranges( + allowed_destination_ranges, reject_overlap=False + ) + except RuntimeError as error: + raise ValueError(str(error)) from error + if actual != expected: + raise ValueError( + "receiver-derived writer plans do not exactly cover the request's KV slots" + ) + return normalized + def writer_base(self, rid_slice: RidSlice, writer_index: int) -> Optional[int]: """Where the given fan-in writer writes in the region.""" with self._reserved_map_lock: ctx = self._reserved_map.get(rid_slice) return None if ctx is None else ctx.writer_base(writer_index) + def bind_writer(self, rid_slice: RidSlice, peer_rank: int, writer_index: int) -> Optional[int]: + """Bind a rank before its sub-region address is advertised.""" + with self._reserved_map_lock: + ctx = self._reserved_map.get(rid_slice) + if ctx is None: + return None + src_base = ctx.writer_base(writer_index) + ctx.bind_writer(peer_rank, src_base) + return src_base + + def set_completion_callback( + self, rid_slice: RidSlice, on_settled: Callable[[bool], None] + ) -> None: + """Install unconditional settlement accounting before any writer is advertised.""" + with self._reserved_map_lock: + ctx = self._reserved_map.get(rid_slice) + if ctx is None: + raise RuntimeError(f"bounce callback for unknown reservation {rid_slice}") + ctx.set_completion_callback(on_settled) + def is_bounced(self, rid_slice: RidSlice) -> bool: with self._reserved_map_lock: return rid_slice in self._reserved_map @@ -304,12 +520,25 @@ def release_idle_reservation(self, rid_slice: RidSlice) -> None: self._recv_alloc.release(ctx.slot_id) def orphan_reservation(self, rid_slice: RidSlice) -> None: - """Give up on a reservation whose write may still be in flight (cancel/timeout/lost result). - The write can't be aborted, so quarantine the region (reclaimed later) rather than releasing - or leaking it. Idempotent; a no-op once the transfer has settled.""" + """Retain a reservation whose remote write may still be in flight. + + A time-based quarantine is not a safety proof: an RMA can outlive any + chosen grace period. ``confirm_drained`` or all writer results must + prove quiescence before this slot returns to the allocator. + """ self._apply(rid_slice, lambda ctx: ctx.mark_orphaned()) - def _apply(self, rid_slice: RidSlice, mutate: Callable[[TransferContext], None]) -> None: + def confirm_drained(self, rid_slice: RidSlice) -> None: + """Release an orphan after sender drain ACK proves reuse is safe.""" + self._apply(rid_slice, lambda ctx: ctx.confirm_drained()) + + def _apply( + self, + rid_slice: RidSlice, + mutate: Callable[[TransferContext], None], + *, + require_present: bool = False, + ) -> None: """Mutate the state under the lock, then do what it asks (scatter or settle) with the lock released, never holding it across a CUDA sync, a queue put, or a callback. No-op if the region is already gone.""" @@ -318,6 +547,8 @@ def _apply(self, rid_slice: RidSlice, mutate: Callable[[TransferContext], None]) with self._reserved_map_lock: ctx = self._reserved_map.get(rid_slice) if ctx is None: + if require_present: + raise RuntimeError(f"bounced result for unknown reservation {rid_slice}") return mutate(ctx) if ctx.ready_to_scatter(): @@ -365,18 +596,104 @@ def record_result( the reader never sees completion before the cache is in place.""" def mut(ctx: TransferContext) -> None: + validated_dst_ptrs, validated_sizes = self._validate_scatter_tail( + ctx, peer_rank, dst_ptrs, sizes, src_base + ) if on_done is not None: ctx.on_done = on_done ctx.record_writer_result( - peer_rank, succeeded=True, src_base=src_base, dst_ptrs=dst_ptrs, sizes=sizes + peer_rank, + succeeded=True, + src_base=src_base, + dst_ptrs=validated_dst_ptrs, + sizes=validated_sizes, ) - self._apply(rid_slice, mut) + self._apply(rid_slice, mut, require_present=True) + + def _validate_scatter_tail( + self, ctx: TransferContext, peer_rank: int, dst_ptrs, sizes, src_base + ) -> tuple[np.ndarray, np.ndarray]: + if dst_ptrs is None or sizes is None or src_base is None: + raise RuntimeError("incomplete bounced-result scatter tail") + dst_ptrs = np.asarray(dst_ptrs, dtype=np.int64) + sizes = np.asarray(sizes, dtype=np.int64) + if dst_ptrs.ndim != 1 or sizes.ndim != 1 or dst_ptrs.size != sizes.size: + raise RuntimeError("invalid bounced-result pointer/size arrays") + if np.any(sizes <= 0): + raise RuntimeError("bounced-result fragment sizes must be positive") + expected_src_base = ctx._expected_writer_bases.get(peer_rank) + if expected_src_base is None or src_base != expected_src_base: + raise RuntimeError(f"bounced result source identity mismatch for rank {peer_rank}") + described_bytes = sum(int(size) for size in sizes) + if described_bytes != ctx.per_writer_bytes: + raise RuntimeError( + f"bounced result describes {described_bytes} bytes, expected {ctx.per_writer_bytes}" + ) + fragments = [ + (int(ptr), int(ptr) + int(size)) for ptr, size in zip(dst_ptrs, sizes, strict=True) + ] + for index, (start, end) in enumerate(fragments): + if index and start < fragments[index - 1][0]: + raise RuntimeError( + "bounced-result destination fragments are not in canonical address order" + ) + if index and start < fragments[index - 1][1]: + raise RuntimeError("bounced-result destination fragments overlap or duplicate") + if not any( + valid_start <= start and end <= valid_end + for valid_start, valid_end in ctx.allowed_destination_ranges + ): + raise RuntimeError( + f"bounced-result destination [{start}, {end}) is outside the " + "receiver-owned KV destination plan" + ) + expected_plan = ctx.expected_destination_plans.get(peer_rank) + actual_plan = tuple( + (int(ptr), int(size)) for ptr, size in zip(dst_ptrs, sizes, strict=True) + ) + if ctx.expected_destination_plans and actual_plan != expected_plan: + raise RuntimeError( + f"bounced-result fragments for rank {peer_rank} do not match the exact " + "receiver-derived destination plan" + ) + return dst_ptrs, sizes - def record_failure(self, rid_slice: RidSlice, peer_rank: int) -> None: + @staticmethod + def _coalesce_destination_ranges( + ranges: list[tuple[int, int]] | tuple[tuple[int, int], ...], + *, + reject_overlap: bool, + ) -> tuple[tuple[int, int], ...]: + coalesced: list[tuple[int, int]] = [] + for start, end in sorted(ranges): + if coalesced and start < coalesced[-1][1]: + if reject_overlap: + raise RuntimeError( + "bounced-result destination fragments overlap or duplicate across writers" + ) + raise RuntimeError("receiver-owned KV destination ranges overlap") + if coalesced and start == coalesced[-1][1]: + coalesced[-1] = (coalesced[-1][0], end) + else: + coalesced.append((start, end)) + return tuple(coalesced) + + def record_failure( + self, + rid_slice: RidSlice, + peer_rank: int, + on_done: Optional[Callable[[bool], None]] = None, + ) -> None: """A writer reported failure (it has drained). The region is freed only once every writer has reported, not here.""" - self._apply(rid_slice, lambda ctx: ctx.record_writer_result(peer_rank, succeeded=False)) + + def mut(ctx: TransferContext) -> None: + if on_done is not None: + ctx.on_done = on_done + ctx.record_writer_result(peer_rank, succeeded=False) + + self._apply(rid_slice, mut) def _scatter_loop(self): CUASSERT(cudart.cudaSetDevice(self._device_id)) @@ -413,6 +730,8 @@ def close(self) -> None: self._scatter_q.put(None) if self._scatter_thread.is_alive(): self._scatter_thread.join(timeout=_CLOSE_JOIN_S) + if self._scatter_thread.is_alive(): + raise RuntimeError("KV bounce scatter thread did not exit; retaining registered memory") for d in self._reg_descs: try: self._agent.deregister_memory(d) @@ -436,13 +755,24 @@ def release_send(self, slot_id) -> None: pass def reserve( - self, recv_req, num_writers: int = 1, *, timeout: Optional[float] = _RESERVE_TIMEOUT_S + self, + recv_req, + num_writers: int = 1, + *, + timeout: Optional[float] = _RESERVE_TIMEOUT_S, + expected_destination_plans: Optional[dict[int, tuple[np.ndarray, np.ndarray]]] = None, ) -> bool: return False def writer_base(self, rid_slice, writer_index: int): return None + def bind_writer(self, rid_slice, peer_rank: int, writer_index: int): + return None + + def set_completion_callback(self, rid_slice, on_settled) -> None: + pass + def is_bounced(self, rid_slice) -> bool: return False @@ -452,13 +782,17 @@ def release_idle_reservation(self, rid_slice) -> None: def orphan_reservation(self, rid_slice) -> None: pass + def confirm_drained(self, rid_slice) -> None: + pass + def record_result( self, rid_slice, peer_rank, dst_ptrs=None, sizes=None, src_base=None, on_done=None ): pass - def record_failure(self, rid_slice, peer_rank) -> None: - pass + def record_failure(self, rid_slice, peer_rank, on_done=None) -> None: + if on_done is not None: + on_done(False) def close(self) -> None: pass @@ -470,8 +804,14 @@ def create_bounce(agent, cfg, *, device_id: int, page_table) -> BounceTransport: if cfg is None: return NoBounceTransport() try: + destination_pool_layouts = _destination_pool_layouts(page_table) transport = VmmBounceTransport.from_config( - agent, cfg, device_id=device_id, block_bytes_per_group=block_bytes_per_group(page_table) + agent, + cfg, + device_id=device_id, + block_bytes_per_group=block_bytes_per_group(page_table), + destination_pool_layouts=destination_pool_layouts, + valid_destination_ranges=[], ) return transport if transport is not None else NoBounceTransport() except ( @@ -488,6 +828,12 @@ def build_send_request(bounce, write_meta, fallback): built = bounce.build_request(write_meta) if built is not None: return built + # The receiver already advertised a leased bounce destination. A + # sender-side in-place fallback would make a tail-less SUCCESS + # indistinguishable from a malformed or stale bounced result. Fail the + # operation explicitly so the receiver retains/retires the reservation + # through the normal terminal protocol. + raise RuntimeError("receiver-advertised bounce request could not be built") return fallback(), None @@ -513,14 +859,14 @@ def encode_result_tail(write_meta) -> list: ] -def decode_result_tail(message): +def decode_result_tail(message, *, tail_index: int = 2): """Recover the destination fragments, sizes, and source from the optional trailing frames, or nothing if the tail is absent.""" - if len(message) >= 5: + if len(message) >= tail_index + 3: return ( - np.frombuffer(message[2], dtype=np.int64), - np.frombuffer(message[3], dtype=np.int64), - int(np.frombuffer(message[4], dtype=np.int64)[0]), + np.frombuffer(message[tail_index], dtype=np.int64), + np.frombuffer(message[tail_index + 1], dtype=np.int64), + int(np.frombuffer(message[tail_index + 2], dtype=np.int64)[0]), ) return None, None, None @@ -528,13 +874,31 @@ def decode_result_tail(message): def block_bytes_per_group(page_table) -> list: """Byte size of one cache block for each leading attention layer group, stopping at the first non-attention group.""" + assert page_table is not None + return [ + sum(slot_bytes for _base, slot_bytes, _slots in group) + for group in _destination_pool_layouts(page_table) + ] + + +def _destination_pool_layouts(page_table) -> list[list[tuple[int, int, int]]]: + """Return deduplicated physical slot layouts for each leading attention group.""" from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool assert page_table is not None - out: list = [] - for lg_idx, lg in enumerate(page_table.layer_groups): - if not isinstance(lg, AttentionLayerGroup): + groups: list[list[tuple[int, int, int]]] = [] + for lg_idx, layer_group in enumerate(page_table.layer_groups): + if not isinstance(layer_group, AttentionLayerGroup): break - out.append(int(get_physical_pool(page_table, lg_idx, 0).slot_bytes)) - return out + unique_layouts = { + ( + int(pool.base_address), + int(pool.slot_bytes), + int(pool.num_slots), + ) + for pool_view in layer_group.pool_views + for pool in [get_physical_pool(page_table, lg_idx, pool_view.pool_idx)] + } + groups.append(sorted(unique_layouts)) + return groups diff --git a/tensorrt_llm/_torch/disaggregation/native/messenger.py b/tensorrt_llm/_torch/disaggregation/native/messenger.py index ceb6aa626ed9..51ad825a24e3 100644 --- a/tensorrt_llm/_torch/disaggregation/native/messenger.py +++ b/tensorrt_llm/_torch/disaggregation/native/messenger.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from abc import ABC, abstractmethod from threading import Event, Lock, Thread from typing import Callable, Optional @@ -190,12 +193,15 @@ def _close_socket(socket: zmq.Socket) -> None: logger.debug("Stopping ZMQMessenger...") self._stop_event.set() - self._internal_socket.send(b"STOP") if self._listener_thread: self._internal_socket.send(b"STOP") self._listener_thread.join(timeout) if self._listener_thread.is_alive(): - logger.warning("Listener thread did not terminate within timeout") + # Closing a socket/context while its listener is still in + # poll/recv is a use-after-close. Keep ownership intact and + # let the caller stop memory teardown. + self._closed = False + raise RuntimeError("Messenger listener did not terminate within timeout") _close_socket(self._socket) _close_socket(self._internal_socket) diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 353444eeca6f..58bbae54430c 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -17,13 +17,15 @@ import os import queue +import secrets import struct import threading import time import weakref -from dataclasses import dataclass +from collections import OrderedDict +from dataclasses import dataclass, field from enum import Enum -from typing import TYPE_CHECKING, List, Optional, Union +from typing import TYPE_CHECKING, Callable, List, Optional, Union import msgpack import numpy as np @@ -55,7 +57,7 @@ WaitResult, ) from tensorrt_llm._torch.disaggregation.native.auxiliary import AuxBuffer -from tensorrt_llm._torch.disaggregation.native.messenger import ZMQMessenger, decode_message +from tensorrt_llm._torch.disaggregation.native.messenger import ZMQMessenger from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfTimer, perf_log_manager @@ -79,6 +81,227 @@ # Number of worker threads for KV transfer queues (default: 1) KV_TRANSFER_NUM_THREADS = int(os.environ.get("TRTLLM_KV_TRANSFER_NUM_THREADS", "1")) +_ASYNC_CONSENSUS_ENVS = ( + "TRTLLM_PYTHON_TRANSCEIVER_ASYNC_CTX_TERMINAL_CONSENSUS", + "TRTLLM_PYTHON_TRANSCEIVER_ASYNC_CTX_PEER_READY_CONSENSUS", +) +_NATIVE_PROTOCOL_VERSION = 2 +_NATIVE_CAPABILITY_PREFIX = b"TRTLLM_NATIVE_TRANSFER_CAPABILITIES\0" +_CONTROL_RETRY_INTERVAL_S = 0.01 +_CONTROL_QUEUE_LIMIT = 65_536 +_LIVE_PROTOCOL_STATE_LIMIT = 65_536 + + +@dataclass(frozen=True) +class _NativeProtocolCapabilities: + version: int = 1 + drain_ack: bool = False + + +_LOCAL_PROTOCOL_CAPABILITIES = _NativeProtocolCapabilities( + version=_NATIVE_PROTOCOL_VERSION, + drain_ack=True, +) +_LEGACY_PROTOCOL_CAPABILITIES = _NativeProtocolCapabilities() + + +def _encode_protocol_capabilities() -> bytes: + return _NATIVE_CAPABILITY_PREFIX + msgpack.packb( + { + "version": _LOCAL_PROTOCOL_CAPABILITIES.version, + "drain_ack": _LOCAL_PROTOCOL_CAPABILITIES.drain_ack, + } + ) + + +def _decode_protocol_capabilities( + frame: Optional[bytes], +) -> _NativeProtocolCapabilities: + if frame is None: + return _LEGACY_PROTOCOL_CAPABILITIES + if not frame.startswith(_NATIVE_CAPABILITY_PREFIX): + raise RuntimeError("invalid native-transfer capability frame") + values = msgpack.unpackb(frame[len(_NATIVE_CAPABILITY_PREFIX) :], raw=False) + return _NativeProtocolCapabilities( + version=int(values.get("version", 1)), + drain_ack=bool(values.get("drain_ack", False)), + ) + + +def _requires_drain_ack_protocol() -> bool: + return any(os.environ.get(name, "0") == "1" for name in _ASYNC_CONSENSUS_ENVS) + + +def _supports_native_protocol_v2(capabilities: _NativeProtocolCapabilities) -> bool: + """Return whether a peer implements the exact protocol used by this process. + + ``drain_ack`` alone is not sufficient: a future peer may retain the bit + while changing the request-incarnation wire contract. Treat unknown + versions as incompatible instead of silently mixing lifetime semantics. + """ + return capabilities.version == _NATIVE_PROTOCOL_VERSION and capabilities.drain_ack + + +@dataclass +class _ControlSend: + endpoint: str + message: list[bytes] + retry: bool + on_sent: Optional[Callable[[], None]] = None + repeat_until: Optional[Callable[[], bool]] = None + done: threading.Event = field(default_factory=threading.Event) + error: Optional[Exception] = None + callback_ran: bool = False + + +class _ControlPlane: + """Own every control DEALER on one thread and serialize its message stream. + + ZMQ sockets are thread-affine. Callers only enqueue immutable frame lists; + the owner thread creates, uses, and closes every socket. Durable messages + retry local send failures until accepted by ZeroMQ, whose reconnect queue + then owns delivery while this progress thread remains alive. + """ + + def __init__(self, name: str): + self._name = name + self._queue: queue.Queue[Optional[_ControlSend]] = queue.Queue(maxsize=_CONTROL_QUEUE_LIMIT) + self._accepting = True + self._state_lock = threading.Lock() + self._owner_ident: Optional[int] = None + self._thread = threading.Thread( + target=self._run, + name=f"trtllm-{name}-control", + daemon=True, + ) + self._thread.start() + + @property + def owner_ident(self) -> Optional[int]: + return self._owner_ident + + def send( + self, + endpoint: Optional[str], + message: list[bytes], + *, + retry: bool = False, + wait: bool = True, + on_sent: Optional[Callable[[], None]] = None, + repeat_until: Optional[Callable[[], bool]] = None, + ) -> _ControlSend: + if endpoint is None: + raise ValueError("control peer endpoint is None; peer may not have registered yet") + request = _ControlSend( + endpoint, + list(message), + retry, + on_sent=on_sent, + repeat_until=repeat_until, + ) + with self._state_lock: + if not self._accepting: + raise RuntimeError(f"{self._name} control plane is shutting down") + try: + self._queue.put_nowait(request) + except queue.Full as error: + raise RuntimeError( + f"{self._name} control queue reached its safety limit " + f"({_CONTROL_QUEUE_LIMIT}); rejecting new control work" + ) from error + if wait: + request.done.wait() + if request.error is not None: + raise RuntimeError( + f"{self._name} control send to {endpoint} failed" + ) from request.error + return request + + def flush(self) -> None: + self._queue.join() + + def shutdown(self) -> None: + with self._state_lock: + if not self._accepting: + should_join = True + else: + self._accepting = False + self._queue.put(None) + should_join = True + if should_join and threading.current_thread() is not self._thread: + self._thread.join() + + def _run(self) -> None: + self._owner_ident = threading.get_ident() + dealers: dict[str, ZMQMessenger] = {} + pending: list[tuple[float, _ControlSend]] = [] + try: + while True: + now = time.monotonic() + ready_index = next( + (index for index, (deadline, _) in enumerate(pending) if deadline <= now), + None, + ) + if ready_index is not None: + _, request = pending.pop(ready_index) + else: + timeout = None + if pending: + timeout = max(0.0, min(deadline for deadline, _ in pending) - now) + try: + request = self._queue.get(timeout=timeout) + except queue.Empty: + continue + if request is None: + self._queue.task_done() + break + try: + if request.repeat_until is not None and request.repeat_until(): + if not request.done.is_set(): + request.done.set() + self._queue.task_done() + continue + dealer = dealers.get(request.endpoint) + if dealer is None: + dealer = ZMQMessenger(mode="DEALER", endpoint=request.endpoint) + dealers[request.endpoint] = dealer + dealer.send(request.message) + if request.on_sent is not None and not request.callback_ran: + request.on_sent() + request.callback_ran = True + if not request.done.is_set(): + request.done.set() + if request.repeat_until is not None and not request.repeat_until(): + pending.append((time.monotonic() + _CONTROL_RETRY_INTERVAL_S, request)) + else: + self._queue.task_done() + except Exception as error: + dealer = dealers.pop(request.endpoint, None) + if dealer is not None: + try: + dealer.stop() + except Exception: + pass + if request.retry: + logger.warning( + "%s control send to %s failed; retrying without " + "blocking unrelated control work: %s", + self._name, + request.endpoint, + error, + ) + pending.append((time.monotonic() + _CONTROL_RETRY_INTERVAL_S, request)) + else: + request.error = error + request.done.set() + self._queue.task_done() + finally: + for dealer in dealers.values(): + try: + dealer.stop() + except Exception as error: + logger.warning("%s control DEALER shutdown failed: %s", self._name, error) + @dataclass class RecvReqInfo: @@ -96,24 +319,28 @@ class RecvReqInfo: mamba_state_index: Optional[int] = None slice_id: Optional[int] = None bounce_dst_base: Optional[int] = None + request_epoch: Optional[int] = None def to_bytes(self) -> bytes: - return msgpack.packb( - { - "sender_req_id": self.sender_req_id, - "instance_name": self.instance_name, - "instance_rank": self.instance_rank, - "block_ids_per_layer_groups": [ - arr.tobytes() for arr in self.block_ids_per_layer_groups - ], - "unique_rid": self.unique_rid, - "dst_start_token": self.dst_start_token, - "aux_slot": self.aux_slot, - "mamba_state_index": self.mamba_state_index, - "slice_id": self.slice_id, - "bounce_dst_base": self.bounce_dst_base, - } - ) + values = { + "sender_req_id": self.sender_req_id, + "instance_name": self.instance_name, + "instance_rank": self.instance_rank, + "block_ids_per_layer_groups": [ + arr.tobytes() for arr in self.block_ids_per_layer_groups + ], + "unique_rid": self.unique_rid, + "dst_start_token": self.dst_start_token, + "aux_slot": self.aux_slot, + "mamba_state_index": self.mamba_state_index, + "slice_id": self.slice_id, + "bounce_dst_base": self.bounce_dst_base, + } + # Keep legacy peers byte-compatible. An epoch is emitted only after + # capability negotiation proves that every writer understands v2. + if self.request_epoch is not None: + values["request_epoch"] = self.request_epoch + return msgpack.packb(values) @classmethod def from_bytes(cls, data: bytes) -> "RecvReqInfo": @@ -152,6 +379,14 @@ class WriteMeta: is_last_slice: bool = False meta_type: WriteMetaType = WriteMetaType.KV bounce_dst_base: Optional[int] = None + request_epoch: Optional[int] = None + # Strongly retain the session until this queued response obligation has + # sent a terminal result to the peer and retired its lifetime credit. + operation_owner: Optional["TxSession"] = None + terminal_sent: bool = False + terminal_queued: bool = False + operation_retired: bool = False + lifetime_lock: threading.Lock = field(default_factory=threading.Lock, repr=False) class MessageType: @@ -162,6 +397,7 @@ class MessageType: REGISTER_RANK_INFO = b"REGISTER_RANK_INFO" AUX_AGENT_RESULT = b"AUX_AGENT_RESULT" CANCEL_SESSION = b"CANCEL_SESSION" + CANCEL_SESSION_ACK = b"CANCEL_SESSION_ACK" class TaskStatus(Enum): @@ -177,35 +413,67 @@ class AgentResult(Enum): # KV_AGENT_RESULT prefix in one struct frame (was ascii frames serialized/parsed under the -# GIL per slice per writer): instance_rank, unique_rid, slice_id, is_last, status, -# transfer_size. The optional bounce tail follows at message[2:]. +# GIL per slice per writer): instance_rank, unique_rid, optional request_epoch, slice_id, +# is_last, status, transfer_size. The optional sender endpoint and bounce tail follow. _KV_RESULT_PREFIX = struct.Struct(" insert time self._peer_requests_lock = threading.Lock() + # Generation-first CTX requests may wait for peer metadata longer + # than the orphan TTL before a TxSession exists. A lease retains only + # their small RecvReqInfo map; no model, KV, or slot resource is held. + self._peer_request_leases: set[int] = set() self._messenger = ZMQMessenger(mode="ROUTER") - self._dealers = {} # used by listener thread only (single-threaded path) + self._control = _ControlPlane("native-sender") self._thread_local = threading.local() # per-thread DEALER cache for worker threads self._sessions = {} # unique_rid -> TxSession - self._sessions_lock = threading.Lock() # Protects _sessions and _pre_cancelled_rids - self._pre_cancelled_rids: set[int] = set() + self._sessions_lock = threading.Lock() + # An acknowledged pre-cancel is live protocol state, not bounded + # history: evicting it can let a delayed session write into memory the + # receiver reclaimed after our ACK. + self._pre_cancelled_rids: dict[int, None] = {} + # Authoritatively-aborted metadata-only requests cannot create a later + # local TxSession. Retain only bounded history to reject in-flight + # setup/REQUEST_DATA races; after eviction, orphan metadata remains + # harmless and is reclaimed by sweep_stale_req_infos(). + self._metadata_cancelled_rids: OrderedDict[int, None] = OrderedDict() + self._pre_cancelled_operations: dict[tuple[int, str, int], None] = {} + self._cancelled_operation_tombstones: OrderedDict[tuple[int, str, int], None] = ( + OrderedDict() + ) + self._closed_rids: OrderedDict[int, None] = OrderedDict() + self._peer_capabilities: dict[str, _NativeProtocolCapabilities] = {} self._shutdown = False + self._ingress_lock = threading.Lock() self._instance_rank = self._registrar.self_rank_info.instance_rank # Guards concurrent add() from the listener thread. self._loaded_remote_agents: set[str] = set() self._loaded_remote_agents_lock = threading.Lock() self._num_threads = KV_TRANSFER_NUM_THREADS self._send_task_queues: List[queue.Queue] = [ - queue.Queue() for _ in range(self._num_threads) + queue.Queue(maxsize=_CONTROL_QUEUE_LIMIT) for _ in range(self._num_threads) ] + self._stalled_operations: list[WriteMeta] = [] + self._stalled_session_owners: list["TxSession"] = [] + self._ambiguous_operations: list[WriteMeta] = [] + self._stalled_operations_lock = threading.Lock() + self._protocol_error: Optional[RuntimeError] = None self._worker_threads: List[threading.Thread] = [ threading.Thread(target=self._process_task_queue, args=(i,), daemon=True) for i in range(self._num_threads) @@ -327,6 +620,15 @@ def _add_req_info(self, unique_rid: int, instance_rank: int, req_info: RecvReqIn if unique_rid not in self._peer_requests: self._peer_requests[unique_rid] = {} self._peer_requests_timestamps[unique_rid] = time.monotonic() + existing = self._peer_requests[unique_rid].get(instance_rank) + if existing is not None and existing.request_epoch != req_info.request_epoch: + error = RuntimeError( + "native-transfer request incarnation changed for an active " + f"operation: rid={unique_rid} receiver_rank={instance_rank} " + f"old_epoch={existing.request_epoch} new_epoch={req_info.request_epoch}" + ) + self._protocol_error = error + raise error self._peer_requests[unique_rid][instance_rank] = req_info def _is_req_ready(self, unique_rid: int, expected_count: int) -> bool: @@ -351,6 +653,21 @@ def _remove_req_info(self, unique_rid: int): with self._peer_requests_lock: self._peer_requests.pop(unique_rid, None) self._peer_requests_timestamps.pop(unique_rid, None) + getattr(self, "_peer_request_leases", set()).discard(unique_rid) + + def pin_peer_req_infos(self, unique_rid: int) -> None: + """Prevent metadata for a known local request from orphan eviction.""" + with self._peer_requests_lock: + if ( + unique_rid not in self._peer_request_leases + and len(self._peer_request_leases) >= _LIVE_PROTOCOL_STATE_LIMIT + ): + self._protocol_error = RuntimeError( + "Sender metadata leases reached their safety limit; " + "refusing to evict live generation-first requests" + ) + raise self._protocol_error + self._peer_request_leases.add(unique_rid) def sweep_stale_req_infos(self): """Evict RecvReqInfo entries that have no matching TxSession and exceed the TTL. @@ -365,12 +682,17 @@ def sweep_stale_req_infos(self): rid for rid, ts in self._peer_requests_timestamps.items() if now - ts > self._STALE_REQ_INFO_TTL_S + and rid not in getattr(self, "_peer_request_leases", set()) ] if not stale_rids: return for rid in stale_rids: with self._sessions_lock, self._peer_requests_lock: - if rid not in self._sessions and rid in self._peer_requests: + if ( + rid not in self._sessions + and rid in self._peer_requests + and rid not in getattr(self, "_peer_request_leases", set()) + ): self._peer_requests.pop(rid, None) self._peer_requests_timestamps.pop(rid, None) logger.debug(f"Swept stale RecvReqInfo for rid={rid}") @@ -378,16 +700,62 @@ def sweep_stale_req_infos(self): def setup_session(self, tx_session: "TxSession"): unique_rid = tx_session.disagg_request_id pre_cancel = False - with self._sessions_lock: - self._sessions[unique_rid] = weakref.ref(tx_session) - if unique_rid in self._pre_cancelled_rids: - pre_cancel = True - self._pre_cancelled_rids.discard(unique_rid) + with self._ingress_lock: + if self._shutdown: + raise RuntimeError("Cannot create a TxSession after Sender shutdown started") + if self._protocol_error is not None: + raise RuntimeError( + "Sender is in fail-stop protocol state" + ) from self._protocol_error + with self._sessions_lock: + self._sessions[unique_rid] = weakref.ref(tx_session) + with self._peer_requests_lock: + getattr(self, "_peer_request_leases", set()).discard(unique_rid) + metadata_cancelled_rids = getattr(self, "_metadata_cancelled_rids", OrderedDict()) + if unique_rid in self._pre_cancelled_rids or unique_rid in metadata_cancelled_rids: + pre_cancel = True + self._pre_cancelled_rids.pop(unique_rid, None) + metadata_cancelled_rids.pop(unique_rid, None) + self._closed_rids.pop(unique_rid, None) if pre_cancel: tx_session.cancel() return - req_info = self._get_first_req_info(unique_rid) + cancelled_operations: list[tuple[str, int]] = [] + with self._ingress_lock, self._sessions_lock: + req_info_map = dict(self._get_req_info(unique_rid) or {}) + for req_info_item in req_info_map.values(): + peer_ri_item = self._registrar.get_peer_rank_info( + req_info_item.instance_name, req_info_item.instance_rank + ) + tx_session.register_request_operation( + peer_ri_item.self_endpoint, req_info_item.request_epoch + ) + if req_info_item.request_epoch is not None: + operation = ( + unique_rid, + peer_ri_item.self_endpoint, + req_info_item.request_epoch, + ) + if operation in self._pre_cancelled_operations: + self._pre_cancelled_operations.pop(operation, None) + self._remember_cancelled_operation_unlocked(operation) + cancelled_operations.append( + (peer_ri_item.self_endpoint, req_info_item.request_epoch) + ) + + if cancelled_operations: + # A request can fan out to multiple receiver endpoints. Preserve + # every exact pre-cancel obligation; acknowledging only the last + # endpoint would strand the other receiver's drain wait forever. + for endpoint, request_epoch in cancelled_operations: + tx_session.cancel( + ack_endpoint=endpoint, + request_epoch=request_epoch, + ) + return + + req_info = next(iter(req_info_map.values()), None) if req_info: peer_ri = self._registrar.get_peer_rank_info( @@ -399,6 +767,45 @@ def setup_session(self, tx_session: "TxSession"): tx_session.receiver_ready = True return + @classmethod + def _remember_tombstone(cls, tombstones: OrderedDict, key) -> None: + tombstones.pop(key, None) + tombstones[key] = None + while len(tombstones) > cls._TOMBSTONE_LIMIT: + tombstones.popitem(last=False) + + def _remember_pre_cancelled_unlocked(self, unique_rid: int) -> None: + if ( + unique_rid not in self._pre_cancelled_rids + and len(self._pre_cancelled_rids) >= _LIVE_PROTOCOL_STATE_LIMIT + ): + self._protocol_error = RuntimeError( + "Sender pre-cancel state reached its safety limit; refusing to evict live state" + ) + raise self._protocol_error + self._pre_cancelled_rids[unique_rid] = None + + def _remember_pre_cancelled_operation_unlocked( + self, unique_rid: int, receiver_endpoint: str, request_epoch: int + ) -> None: + key = (unique_rid, receiver_endpoint, request_epoch) + if ( + key not in self._pre_cancelled_operations + and len(self._pre_cancelled_operations) >= _LIVE_PROTOCOL_STATE_LIMIT + ): + self._protocol_error = RuntimeError( + "Sender v2 pre-cancel state reached its safety limit; " + "refusing to evict live request incarnations" + ) + raise self._protocol_error + self._pre_cancelled_operations[key] = None + + def _remember_closed_unlocked(self, unique_rid: int) -> None: + self._remember_tombstone(self._closed_rids, unique_rid) + + def _remember_cancelled_operation_unlocked(self, operation: tuple[int, str, int]) -> None: + self._remember_tombstone(self._cancelled_operation_tombstones, operation) + def _get_session(self, unique_rid: Optional[int]) -> Optional["TxSession"]: session_ref = self._sessions.get(unique_rid) if session_ref is None: @@ -414,7 +821,16 @@ def _enqueue(self, write_meta: WriteMeta): # - Same peer's slices stay ordered on one thread (is_last_slice correctness) # - Different peers can run on different threads (better load balancing) thread_idx = hash((write_meta.unique_rid, write_meta.peer_rank)) % self._num_threads - self._send_task_queues[thread_idx].put(write_meta) + with self._ingress_lock: + if self._shutdown: + raise RuntimeError("Sender is shutting down; transfer enqueue rejected") + try: + self._send_task_queues[thread_idx].put_nowait(write_meta) + except queue.Full as error: + raise RuntimeError( + "native transfer work queue reached its safety limit; " + "rejecting new transfer work" + ) from error def _get_or_connect_thread_dealer(self, endpoint: Optional[str]) -> ZMQMessenger: """Get or create a per-thread DEALER socket via threading.local(). @@ -430,6 +846,99 @@ def _get_or_connect_thread_dealer(self, endpoint: Optional[str]) -> ZMQMessenger dealers[endpoint] = ZMQMessenger(mode="DEALER", endpoint=endpoint) return dealers[endpoint] + def _send_write_result( + self, + write_meta: WriteMeta, + result: AgentResult, + *, + transfer_size: int = 0, + tail: Optional[list] = None, + ) -> None: + """Publish exactly one terminal response from the owning worker thread.""" + if write_meta.terminal_sent or write_meta.terminal_queued: + return + if write_meta.meta_type == WriteMetaType.KV: + if write_meta.slice_id is None: + raise RuntimeError("KV WriteMeta is missing slice_id") + message = _make_kv_result_msg( + self._instance_rank, + write_meta.unique_rid, + write_meta.slice_id, + write_meta.is_last_slice if result == AgentResult.SUCCESS else True, + result, + transfer_size=transfer_size, + tail=tail, + request_epoch=write_meta.request_epoch, + sender_endpoint=self.endpoint, + ) + else: + message = [ + MessageType.AUX_AGENT_RESULT, + str(self._instance_rank).encode("ascii"), + str(write_meta.unique_rid).encode("ascii"), + result.value.encode("ascii"), + ] + if write_meta.request_epoch is not None: + message.extend( + [ + self.endpoint.encode("utf-8"), + str(write_meta.request_epoch).encode("ascii"), + ] + ) + try: + self._get_or_connect_thread_dealer(write_meta.peer_endpoint).send(message) + write_meta.terminal_sent = True + except Exception as error: + # The RMA is terminal but the receiver still owns one lifetime + # credit. Move only this rare error path to the durable control + # owner; its callback holds the session strongly and retires the + # credit only after ZeroMQ accepts the terminal notification. + logger.warning( + "Terminal result send failed for request %s peer_rank=%s; " + "moving it to the durable control queue: %s", + write_meta.unique_rid, + write_meta.peer_rank, + error, + ) + write_meta.terminal_queued = True + + def terminal_sent() -> None: + write_meta.terminal_sent = True + self._retire_write_meta_operation(write_meta) + with self._stalled_operations_lock: + if write_meta in self._stalled_operations: + self._stalled_operations.remove(write_meta) + + try: + self._control.send( + write_meta.peer_endpoint, + message, + retry=True, + wait=False, + on_sent=terminal_sent, + ) + except Exception: + write_meta.terminal_queued = False + raise + + def _fail_write_meta(self, write_meta: WriteMeta, error: Exception) -> None: + write_meta.task.fail(error) + owner = write_meta.operation_owner + if owner is not None: + owner.set_exception(str(error)) + self._send_write_result(write_meta, AgentResult.FAILED) + + @staticmethod + def _retire_write_meta_operation(write_meta: WriteMeta) -> None: + owner = write_meta.operation_owner + if owner is None: + return + with write_meta.lifetime_lock: + if write_meta.operation_retired: + return + owner.retire_operation() + write_meta.operation_retired = True + def _process_task_queue(self, thread_idx: int): device_id = self._device_id torch.cuda.set_device(device_id) @@ -449,12 +958,63 @@ def _process_task_queue(self, thread_idx: int): self._deliver_aux_to_agent(write_meta) else: self._deliver_kv_to_agent(write_meta) + except _AmbiguousTransferError as error: + # submit()/wait() raised after the agent may have accepted + # the RMA. Publishing FAILED or retiring the lifetime + # credit would falsely claim drain. Retain everything and + # force teardown to stop before memory invalidation. + logger.critical( + "Ambiguous transfer state for request %s: %s", + write_meta.unique_rid, + error, + ) + write_meta.task.fail(error) + owner = write_meta.operation_owner + if owner is not None: + owner.set_exception(str(error)) + with self._stalled_operations_lock: + if write_meta not in self._ambiguous_operations: + self._ambiguous_operations.append(write_meta) + continue except Exception as e: logger.error( f"_process_task_queue[{thread_idx}]: unhandled exception for " f"unique_rid={write_meta.unique_rid}: {e}" ) - write_meta.task.fail(e) + try: + self._fail_write_meta(write_meta, e) + except Exception as notify_error: + logger.error( + "Unable to publish terminal failure for request %s: %s", + write_meta.unique_rid, + notify_error, + ) + finally: + owner = write_meta.operation_owner + if owner is not None and write_meta.terminal_sent: + try: + self._retire_write_meta_operation(write_meta) + except Exception as retire_error: + logger.error( + "Unable to retire terminal operation for request %s: %s", + write_meta.unique_rid, + retire_error, + ) + with self._stalled_operations_lock: + if write_meta not in self._stalled_operations: + self._stalled_operations.append(write_meta) + elif owner is not None: + # Serialize with the durable callback's removal. It may + # have published the terminal result between the first + # flag check above and this lock acquisition. + with self._stalled_operations_lock: + if ( + not write_meta.terminal_sent + and write_meta not in self._stalled_operations + ): + # No terminal notification has retired this + # credit. Retain it strongly and fail closed. + self._stalled_operations.append(write_meta) finally: # Clean up this thread's DEALER sockets. threading.local storage # is only accessible from the owning thread, so shutdown must @@ -510,56 +1070,69 @@ def _make_agent_request(write_meta: WriteMeta, device_id: int) -> "TransferReque TransferOp.WRITE, src_memory_descs, dst_memory_descs, write_meta.peer_name, None ) + def _submit_and_wait(self, request, write_meta: WriteMeta) -> tuple[bool, object]: + """Return a known terminal result or raise without claiming quiescence. + + Either submit or wait may throw after the agent accepted the request. + Such an exception is ambiguous by contract: no terminal peer result is + sent and registered memory remains retained through fail-stop teardown. + """ + try: + status = self._agent.submit_transfer_requests(request) + except Exception as error: + raise _AmbiguousTransferError( + f"transfer submission may have been accepted for request " + f"{write_meta.unique_rid} peer_rank={write_meta.peer_rank}" + ) from error + try: + return bool(status.wait()), status + except Exception as error: + raise _AmbiguousTransferError( + f"transfer completion is unknown for request {write_meta.unique_rid} " + f"peer_rank={write_meta.peer_rank}" + ) from error + @nvtx_range("_deliver_kv_to_agent") def _deliver_kv_to_agent(self, write_meta: WriteMeta): assert write_meta.src_ptrs.size == write_meta.dst_ptrs.size == write_meta.sizes.size, ( f"WriteMeta ptr/size mismatch for unique_rid={write_meta.unique_rid}" ) - with self._sessions_lock: - session = self._get_session(write_meta.unique_rid) + session = write_meta.operation_owner + if session is None: + with self._sessions_lock: + session = self._get_session(write_meta.unique_rid) if session is None: msg = ( f"_deliver_kv_to_agent: TxSession {write_meta.unique_rid} not found or already GC'd" ) logger.error(msg) write_meta.task.fail(RuntimeError(msg)) + self._send_write_result(write_meta, AgentResult.FAILED) return assert write_meta.slice_id is not None task = session.kv_tasks[write_meta.slice_id] timer = task._perf_timer if timer: timer.record_push_end(write_meta.peer_rank) - # Hold session.lock to serialize the INIT→TRANSFERRING transition with - # cancel(): prevents cancel_request() from freeing KV pages while a - # worker is about to write into them. - with session.lock: + # The session owns the INIT-to-TRANSFERRING transition. Its permanent + # seal prevents a queued task from starting after a terminal vote. + if not session.try_mark_transferring(task): status = session.status - if status in (SessionStatus.ERROR, SessionStatus.CANCELLED): - should_abort = True - else: - task.status = TaskStatus.TRANSFERRING - should_abort = False - - if should_abort: logger.warning( f"_deliver_kv_to_agent: session {write_meta.unique_rid} already " - f"in {status.value} state; sending FAILED to receiver" + f"sealed or in {status.value} state; sending FAILED to receiver" ) # Task may have been enqueued after cancel() already iterated kv_tasks, # so its future was never set by cancel(). Set it here as a fallback. - task.fail( - RuntimeError(f"session {write_meta.unique_rid} {status.value}, transfer aborted") - ) - self._get_or_connect_dealer(write_meta.peer_endpoint).send( - _make_kv_result_msg( - self._instance_rank, - write_meta.unique_rid, - write_meta.slice_id, - True, # is_last_slice — ensures receiver resolves its task future - AgentResult.FAILED, + if not task.is_done: + task.fail( + RuntimeError( + f"session {write_meta.unique_rid} sealed or {status.value}, " + "transfer aborted" + ) ) - ) + self._send_write_result(write_meta, AgentResult.FAILED) return from .bounce import build_send_request, encode_result_tail @@ -581,21 +1154,15 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): f"{write_meta.unique_rid} slice={write_meta.slice_id}: {e}" ) task.fail(RuntimeError(f"build_send_request failed: {e}")) - self._get_or_connect_dealer(write_meta.peer_endpoint).send( - _make_kv_result_msg( - self._instance_rank, - write_meta.unique_rid, - write_meta.slice_id, - True, # is_last_slice — ensures receiver resolves its task future - AgentResult.FAILED, - ) - ) + self._send_write_result(write_meta, AgentResult.FAILED) return if timer: timer.record_transfer_start(write_meta.peer_rank) + terminal_known = False try: - status = self._agent.submit_transfer_requests(request) - if not status.wait(): + completed, status = self._submit_and_wait(request, write_meta) + terminal_known = True + if not completed: agent_result = AgentResult.FAILED last_status = getattr(status, "last_status_str", lambda: "")() agent_name = getattr(self._agent, "name", "") @@ -614,7 +1181,7 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): logger.error(detail) task.fail(RuntimeError(detail)) finally: - if send_slot_id is not None: + if send_slot_id is not None and terminal_known: self._bounce.release_send(send_slot_id) if timer: timer.record_transfer_end(write_meta.peer_rank) @@ -626,16 +1193,12 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): else None ) transfer_size = timer.get_transfer_size(write_meta.peer_rank) if timer else 0 - result_msg = _make_kv_result_msg( - self._instance_rank, - write_meta.unique_rid, - write_meta.slice_id, - write_meta.is_last_slice, + self._send_write_result( + write_meta, agent_result, transfer_size=transfer_size, tail=tail, ) - self._get_or_connect_thread_dealer(write_meta.peer_endpoint).send(result_msg) if timer: timer.record_task_end(write_meta.peer_rank) @@ -647,6 +1210,12 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): count = task.transferred_count if count > write_meta.expected_transfers: + task.fail( + RuntimeError( + f"KV slice {write_meta.slice_id} received more than " + f"{write_meta.expected_transfers} transfers" + ) + ) session.set_exception( f"KV slice {write_meta.slice_id} received more than {write_meta.expected_transfers} transfers" ) @@ -668,11 +1237,12 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): @nvtx_range("_deliver_aux_to_agent") def _deliver_aux_to_agent(self, write_meta: WriteMeta): - session = self._get_session(write_meta.unique_rid) + session = write_meta.operation_owner or self._get_session(write_meta.unique_rid) if session is None: msg = f"_deliver_aux_to_agent: TxSession {write_meta.unique_rid} not found or already GC'd" logger.error(msg) write_meta.task.fail(RuntimeError(msg)) + self._send_write_result(write_meta, AgentResult.FAILED) return aux_task = session.aux_task assert aux_task is not None, f"aux_task is None for session {write_meta.unique_rid}" @@ -680,25 +1250,38 @@ def _deliver_aux_to_agent(self, write_meta: WriteMeta): if timer: timer.record_push_end(write_meta.peer_rank) + if not session.try_mark_transferring(aux_task): + status = session.status + logger.warning( + f"_deliver_aux_to_agent: session {write_meta.unique_rid} already " + f"sealed or in {status.value} state; sending FAILED to receiver" + ) + if not aux_task.is_done: + aux_task.fail( + RuntimeError( + f"session {write_meta.unique_rid} sealed or {status.value}, " + "aux transfer aborted" + ) + ) + self._send_write_result(write_meta, AgentResult.FAILED) + return + agent_result = AgentResult.SUCCESS if write_meta.src_ptrs.size > 0: request = Sender._make_agent_request(write_meta, device_id=self._device_id) if timer: timer.record_transfer_start(write_meta.peer_rank) - if not self._agent.submit_transfer_requests(request).wait(): + completed, _status = self._submit_and_wait(request, write_meta) + if not completed: agent_result = AgentResult.FAILED + aux_task.fail( + RuntimeError(f"aux transfer agent request failed for {write_meta.unique_rid}") + ) session.set_exception("aux transfer agent request failed") if timer: timer.record_transfer_end(write_meta.peer_rank) - self._get_or_connect_thread_dealer(write_meta.peer_endpoint).send( - [ - MessageType.AUX_AGENT_RESULT, - str(self._instance_rank).encode("ascii"), - str(write_meta.unique_rid).encode("ascii"), - agent_result.value.encode("ascii"), - ] - ) + self._send_write_result(write_meta, agent_result) if timer: timer.record_task_end(write_meta.peer_rank) @@ -716,6 +1299,11 @@ def _deliver_aux_to_agent(self, write_meta: WriteMeta): else: aux_task.complete() elif count > write_meta.expected_transfers: + aux_task.fail( + RuntimeError( + f"aux task received more than {write_meta.expected_transfers} transfers" + ) + ) session.set_exception( f"aux task received more than {write_meta.expected_transfers} transfers" ) @@ -914,6 +1502,7 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write slice_id=task.slice_id, is_last_slice=task._slice.is_last_slice, bounce_dst_base=req_info.bounce_dst_base, + request_epoch=req_info.request_epoch, ) def _build_aux_write_meta(self, task: AuxSendTask, req_info: RecvReqInfo) -> WriteMeta: @@ -955,27 +1544,78 @@ def _build_aux_write_meta(self, task: AuxSendTask, req_info: RecvReqInfo) -> Wri peer_endpoint=peer_ri.self_endpoint, unique_rid=task._unique_rid, meta_type=WriteMetaType.AUX, + request_epoch=req_info.request_epoch, ) def dispatch_task( self, task: KVSendTask | AuxSendTask, req_info_snapshot: Optional[dict] = None, + operation_owner: Optional["TxSession"] = None, ): # req_info_snapshot may be pre-fetched under session.lock by the caller to keep the # critical section small. When not provided, we fetch it here (legacy / standalone path). if req_info_snapshot is None: req_info_snapshot = dict(self._get_req_info(task._unique_rid) or {}) - for info in req_info_snapshot.values(): - if task._perf_timer is not None: - task._perf_timer.record_task_start(info.instance_rank) - if isinstance(task, KVSendTask): - trans_meta = self._build_kv_write_meta(task, info) - else: - trans_meta = self._build_aux_write_meta(task, info) - if task._perf_timer is not None: - task._perf_timer.record_push_start(trans_meta.peer_rank) - self._enqueue(trans_meta) + write_metas: list[WriteMeta] = [] + try: + for info in req_info_snapshot.values(): + if operation_owner is not None: + peer_ri = self._registrar.get_peer_rank_info( + info.instance_name, info.instance_rank + ) + operation_owner.register_request_operation( + peer_ri.self_endpoint, info.request_epoch + ) + if task._perf_timer is not None: + task._perf_timer.record_task_start(info.instance_rank) + if isinstance(task, KVSendTask): + trans_meta = self._build_kv_write_meta(task, info) + else: + trans_meta = self._build_aux_write_meta(task, info) + if task._perf_timer is not None: + task._perf_timer.record_push_start(trans_meta.peer_rank) + write_metas.append(trans_meta) + except Exception as error: + task.fail(error) + if operation_owner is not None: + operation_owner.set_exception(str(error)) + infos = list(req_info_snapshot.values()) + if operation_owner is not None: + operation_owner.finish_failed_dispatch(len(infos)) + notification_failed = False + for info in infos: + on_sent = operation_owner.retire_operation if operation_owner is not None else None + if isinstance(task, KVSendTask): + queued = self._send_failed_result_to_receiver( + info, retry=True, wait=False, on_sent=on_sent + ) + else: + queued = self._send_failed_aux_result_to_receiver( + info, retry=True, wait=False, on_sent=on_sent + ) + notification_failed = notification_failed or not queued + if operation_owner is not None and notification_failed: + # A terminal notification has no durable queue owner. Keep the + # session strongly reachable and fail closed so its transfer + # memory cannot be reclaimed while the receiver may still be + # waiting or writing. + with self._stalled_operations_lock: + self._stalled_session_owners.append(operation_owner) + return + + if operation_owner is not None: + operation_owner.finish_dispatch(write_metas) + for write_meta in write_metas: + try: + self._enqueue(write_meta) + except Exception as error: + self._fail_write_meta(write_meta, error) + if write_meta.operation_owner is not None and write_meta.terminal_sent: + write_meta.operation_owner.retire_operation() + elif write_meta.operation_owner is not None: + with self._stalled_operations_lock: + self._stalled_operations.append(write_meta) def _start_listener(self): def handle_message(messages: list[bytes]): @@ -1011,8 +1651,10 @@ def _register_peer_rank(self, _send_id: bytes, message: list[bytes]): torch.cuda.set_device(self._device_id) CUASSERT(cudart.cudaSetDevice(self._device_id)) ri: RankInfo = RankInfo.from_bytes(message[1]) + capabilities = _decode_protocol_capabilities(message[2] if len(message) > 2 else None) self._registrar.register(ri.instance_name, ri.instance_rank, ri) + self._peer_capabilities[ri.self_endpoint] = capabilities agent_name = ri.instance_name + str(ri.instance_rank) logger.debug(f"Loading remote transfer agent descriptor for peer '{agent_name}'") @@ -1028,28 +1670,122 @@ def _register_peer_rank(self, _send_id: bytes, message: list[bytes]): def _handle_cancel_session(self, message: list[bytes]): unique_rid = int(message[1]) + ack_endpoint = message[2].decode("utf-8") if len(message) > 2 else None + request_epoch = int(message[3]) if len(message) > 3 else None session = None - with self._sessions_lock: - session_ref = self._sessions.get(unique_rid) - if session_ref is None: - self._pre_cancelled_rids.add(unique_rid) + resend_ack = False + with self._ingress_lock, self._sessions_lock: + if request_epoch is not None: + if ack_endpoint is None: + raise RuntimeError("v2 cancel is missing its receiver endpoint") + operation = (unique_rid, ack_endpoint, request_epoch) + if operation in self._cancelled_operation_tombstones: + resend_ack = True + else: + session_ref = self._sessions.get(unique_rid) + session = session_ref() if session_ref is not None else None + if session is not None and session.has_request_operation( + ack_endpoint, request_epoch + ): + pass + elif session is not None and session.has_request_endpoint(ack_endpoint): + error = RuntimeError( + "stale native-transfer cancel does not match the active " + f"request incarnation: rid={unique_rid} endpoint={ack_endpoint} " + f"epoch={request_epoch}" + ) + self._protocol_error = error + raise error + else: + session = None + self._remember_pre_cancelled_operation_unlocked( + unique_rid, ack_endpoint, request_epoch + ) + resend_ack = True else: - session = session_ref() - if session is None: - self._pre_cancelled_rids.add(unique_rid) + session_ref = self._sessions.get(unique_rid) + if session_ref is None: + if unique_rid not in self._closed_rids: + self._remember_pre_cancelled_unlocked(unique_rid) + else: + session = session_ref() + if session is None and unique_rid not in self._closed_rids: + self._remember_pre_cancelled_unlocked(unique_rid) if session is not None: - session.cancel() + session.cancel(ack_endpoint=ack_endpoint, request_epoch=request_epoch) + elif resend_ack and ack_endpoint is not None: + self.send_cancel_ack( + ack_endpoint, + unique_rid, + request_epoch=request_epoch, + from_worker=False, + ) @nvtx_range("_respond_with_kv") def _respond_with_kv(self, _send_id: bytes, message: list[bytes]): # _sessions_lock prevents a race between session lookup and req_info save. # session.lock serializes _enqueue calls from both paths. info: RecvReqInfo = RecvReqInfo.from_bytes(message[1]) - with self._sessions_lock: + peer_ri = self._registrar.get_peer_rank_info(info.instance_name, info.instance_rank) + if self._shutdown: + self._send_failed_result_to_receiver(info, retry=True, wait=False) + return + capabilities = self._peer_capabilities.get( + peer_ri.self_endpoint, _LEGACY_PROTOCOL_CAPABILITIES + ) + if info.request_epoch is not None and not _supports_native_protocol_v2(capabilities): + logger.error( + "Rejecting v2 request %s from incompatible peer %s (version=%s)", + info.unique_rid, + peer_ri.self_endpoint, + capabilities.version, + ) + self._send_failed_result_to_receiver(info, retry=True, wait=False) + return + if _requires_drain_ack_protocol() and not _supports_native_protocol_v2(capabilities): + logger.error( + "Rejecting request %s from legacy native-transfer peer %s: " + "asynchronous consensus requires drain-ACK protocol version %s", + info.unique_rid, + peer_ri.self_endpoint, + _NATIVE_PROTOCOL_VERSION, + ) + self._send_failed_result_to_receiver(info, retry=True, wait=False) + return + cancelled_before_request = False + with self._ingress_lock, self._sessions_lock: session = self._get_session(info.unique_rid) + if info.request_epoch is not None: + operation = ( + info.unique_rid, + peer_ri.self_endpoint, + info.request_epoch, + ) + if operation in self._pre_cancelled_operations: + self._pre_cancelled_operations.pop(operation, None) + self._remember_cancelled_operation_unlocked(operation) + cancelled_before_request = True + elif operation in self._cancelled_operation_tombstones: + cancelled_before_request = True + elif session is not None: + session.register_request_operation(peer_ri.self_endpoint, info.request_epoch) if session is None: - self._save_peer_req_info(info) - return + if ( + cancelled_before_request + or info.unique_rid in self._closed_rids + or info.unique_rid in self._pre_cancelled_rids + or info.unique_rid in getattr(self, "_metadata_cancelled_rids", OrderedDict()) + ): + cancelled_before_request = True + else: + self._save_peer_req_info(info) + session = None + if cancelled_before_request: + self._send_failed_result_to_receiver(info) + return + if session is None: + return + dispatch_tasks: list[SendTaskBase] = [] with session.lock: self._save_peer_req_info(info) tasks = list(session.kv_tasks) @@ -1059,38 +1795,100 @@ def _respond_with_kv(self, _send_id: bytes, message: list[bytes]): if not tasks and session.status in (SessionStatus.ERROR, SessionStatus.CANCELLED): self._send_failed_result_to_receiver(info) return - for task in tasks: - if task._perf_timer is not None: - task._perf_timer.record_task_start(info.instance_rank) - trans_meta = self._build_kv_write_meta(task, info) - if task._perf_timer is not None: - task._perf_timer.record_push_start(trans_meta.peer_rank) - self._enqueue(trans_meta) - - def _send_failed_result_to_receiver(self, info: RecvReqInfo): + for task in tasks: + if session.begin_dispatch_unlocked(): + dispatch_tasks.append(task) + else: + self._send_failed_result_to_receiver(info) + for task in dispatch_tasks: + self.dispatch_task(task, {info.instance_rank: info}, operation_owner=session) + + def _send_failed_result_to_receiver( + self, + info: RecvReqInfo, + *, + retry: bool = False, + wait: bool = True, + on_sent: Optional[Callable[[], None]] = None, + ) -> bool: try: peer_ri = self._registrar.get_peer_rank_info(info.instance_name, info.instance_rank) slice_id = info.slice_id if info.slice_id is not None else 0 - self._get_or_connect_dealer(peer_ri.self_endpoint).send( + self._send_control_message( + peer_ri.self_endpoint, _make_kv_result_msg( self._instance_rank, info.unique_rid, slice_id, - True, # is_last_slice + True, AgentResult.FAILED, - ) + request_epoch=info.request_epoch, + sender_endpoint=self.endpoint, + ), + retry=retry, + wait=wait, + on_sent=on_sent, ) + return True except Exception as e: logger.warning( f"_respond_with_kv: failed to abort receiver for rid={info.unique_rid}: {e}" ) + return False - def _get_or_connect_dealer(self, endpoint: Optional[str]): - if endpoint is None: - raise ValueError("Sender: peer endpoint is None; peer may not have registered yet") - if endpoint not in self._dealers: - self._dealers[endpoint] = ZMQMessenger(mode="DEALER", endpoint=endpoint) - return self._dealers[endpoint] + def _send_failed_aux_result_to_receiver( + self, + info: RecvReqInfo, + *, + retry: bool = False, + wait: bool = True, + on_sent: Optional[Callable[[], None]] = None, + ) -> bool: + try: + peer_ri = self._registrar.get_peer_rank_info(info.instance_name, info.instance_rank) + message = [ + MessageType.AUX_AGENT_RESULT, + str(self._instance_rank).encode("ascii"), + str(info.unique_rid).encode("ascii"), + AgentResult.FAILED.value.encode("ascii"), + ] + if info.request_epoch is not None: + message.extend( + [ + self.endpoint.encode("utf-8"), + str(info.request_epoch).encode("ascii"), + ] + ) + self._send_control_message( + peer_ri.self_endpoint, + message, + retry=retry, + wait=wait, + on_sent=on_sent, + ) + return True + except Exception as error: + logger.warning(f"Failed to abort aux receiver for rid={info.unique_rid}: {error}") + return False + + def _send_control_message( + self, + endpoint: Optional[str], + message: list[bytes], + *, + retry: bool = False, + wait: bool = True, + on_sent: Optional[Callable[[], None]] = None, + repeat_until: Optional[Callable[[], bool]] = None, + ) -> _ControlSend: + return self._control.send( + endpoint, + message, + retry=retry, + wait=wait, + on_sent=on_sent, + repeat_until=repeat_until, + ) def _save_peer_req_info(self, peer_transfer_req_info: RecvReqInfo): req_info = peer_transfer_req_info @@ -1113,10 +1911,19 @@ def _has_all_peer_req_infos(self, req_info: RecvReqInfo) -> bool: expected_transfers = len(self._registrar.get_peer_overlap(peer_ri, peer_ri.dp_rank).ranks) return self._is_req_ready(req_info.unique_rid, expected_transfers) - def clear_session(self, unique_rid: int): + def clear_session( + self, + unique_rid: int, + drained_operations: Optional[set[tuple[str, Optional[int]]]] = None, + ): with self._sessions_lock: - if unique_rid in self._sessions: - del self._sessions[unique_rid] + self._sessions.pop(unique_rid, None) + self._remember_closed_unlocked(unique_rid) + for endpoint, request_epoch in drained_operations or (): + if request_epoch is not None: + self._remember_cancelled_operation_unlocked( + (unique_rid, endpoint, request_epoch) + ) self._remove_req_info(unique_rid) def send_cancel_to_receivers(self, unique_rid: int) -> None: @@ -1126,29 +1933,124 @@ def send_cancel_to_receivers(self, unique_rid: int) -> None: with self._peer_requests_lock: req_info_map = self._peer_requests.get(unique_rid) req_infos = list(req_info_map.values()) if req_info_map else [] + self._send_cancel_to_req_infos(unique_rid, req_infos) + + def _send_cancel_to_req_infos(self, unique_rid: int, req_infos: list[RecvReqInfo]) -> None: for req_info in req_infos: try: peer_ri = self._registrar.get_peer_rank_info( req_info.instance_name, req_info.instance_rank ) - self._get_or_connect_dealer(peer_ri.self_endpoint).send( - [MessageType.CANCEL_SESSION, str(unique_rid).encode("ascii")] + self._send_control_message( + peer_ri.self_endpoint, + [ + MessageType.CANCEL_SESSION, + str(unique_rid).encode("ascii"), + self.endpoint.encode("utf-8"), + *( + [str(req_info.request_epoch).encode("ascii")] + if req_info.request_epoch is not None + else [] + ), + ], + retry=True, + wait=False, ) except Exception as e: logger.warning(f"send_cancel_to_receivers: failed for rid={unique_rid}: {e}") - def shutdown(self): - if self._shutdown: - return - self._shutdown = True + def cancel_peer_req_infos(self, unique_rid: int) -> bool: + """Retire a metadata-only send request before TxSession creation. - # Quiesce listener before invalidate to avoid set/map mutation races. - self._messenger.stop() + The caller may use this only after an authoritative readiness abort + guarantees that no later local TxSession can be created for the + process-unique request ID. A bounded tombstone makes in-flight + setup/REQUEST_DATA races fail closed. After that history expires, late + ingress can create only orphan metadata, which the stale sweep reclaims. + Returns False if a live TxSession has already assumed ownership. + """ + with self._ingress_lock, self._sessions_lock: + session_ref = self._sessions.get(unique_rid) + if session_ref is not None and session_ref() is not None: + return False + self._sessions.pop(unique_rid, None) + metadata_cancelled_rids = getattr(self, "_metadata_cancelled_rids", None) + if metadata_cancelled_rids is None: + metadata_cancelled_rids = OrderedDict() + self._metadata_cancelled_rids = metadata_cancelled_rids + self._remember_tombstone(metadata_cancelled_rids, unique_rid) + self._remember_closed_unlocked(unique_rid) + with self._peer_requests_lock: + req_info_map = self._peer_requests.pop(unique_rid, None) + req_infos = list(req_info_map.values()) if req_info_map else [] + self._peer_requests_timestamps.pop(unique_rid, None) + self._peer_request_leases.discard(unique_rid) + self._send_cancel_to_req_infos(unique_rid, req_infos) + return True + + def send_cancel_ack( + self, + endpoint: str, + unique_rid: int, + *, + request_epoch: Optional[int] = None, + from_worker: bool, + ) -> bool: + """Acknowledge that all sender-side writes for a cancelled session drained.""" + try: + message = [ + MessageType.CANCEL_SESSION_ACK, + str(unique_rid).encode("ascii"), + self.endpoint.encode("utf-8"), + ] + if request_epoch is not None: + message.append(str(request_epoch).encode("ascii")) + self._send_control_message(endpoint, message, retry=True, wait=False) + return True + except Exception as error: + logger.warning(f"Failed to acknowledge cancel for rid={unique_rid}: {error}") + return False - for q in self._send_task_queues: - q.put(None) + def shutdown(self): + with self._ingress_lock: + if self._shutdown: + return + self._shutdown = True + for q in self._send_task_queues: + q.put(None) for t in self._worker_threads: t.join(timeout=5) + live_workers = [t for t in self._worker_threads if t.is_alive()] + if live_workers: + logger.warning( + f"Sender shutdown is waiting for {len(live_workers)} transfer worker(s) " + "to drain before invalidating remote agents" + ) + # Registered transfer memory and remote-agent descriptors must + # outlive every in-flight worker. Returning after a timed join + # would race memory teardown, so shutdown deliberately fails + # closed and waits for the transfer agent's own bounded wait. + for t in live_workers: + t.join() + with self._stalled_operations_lock: + if self._ambiguous_operations: + raise RuntimeError( + "Sender has transfer-agent operations with unknown terminal state; " + "retaining control sockets, remote agents, and registered memory" + ) + self._control.flush() + # Keep cancellation/ACK progress alive until every worker and queued + # terminal notification has drained. Stop ingress only at that point. + self._messenger.stop() + self._control.flush() + with self._stalled_operations_lock: + if self._stalled_operations or self._stalled_session_owners: + logger.error( + "Sender shutdown retained %s operation(s) and %s session owner(s) " + "whose terminal peer response could not be published", + len(self._stalled_operations), + len(self._stalled_session_owners), + ) # Snapshot under lock as defense in depth. with self._loaded_remote_agents_lock: @@ -1162,12 +2064,7 @@ def shutdown(self): logger.warning( f"Failed to invalidate remote agent '{agent_name}' during shutdown: {e}" ) - for dealer in self._dealers.values(): - try: - dealer.stop() - except Exception as e: - logger.warning(f"Failed to stop dealer during Sender shutdown: {e}") - self._dealers.clear() + self._control.shutdown() def __del__(self): try: @@ -1210,9 +2107,19 @@ def __init__( self._exception: Optional[Exception] = None self._closed = False + self._sealed = False self._terminal_status: Optional[SessionStatus] = None self.transfer_start_time = None self.transfer_end_time = None + self._terminal_snapshot: Optional[SessionStatus] = None + self._outstanding_operations = 0 + self._close_requested = False + self._cancel_ack_endpoints: set[str] = set() + self._cancel_acked_endpoints: set[str] = set() + self._request_operations: set[tuple[str, Optional[int]]] = set() + self._cancel_ack_operations: set[tuple[str, Optional[int]]] = set() + self._cancel_acked_operations: set[tuple[str, Optional[int]]] = set() + self._cancel_notified = False # Must be last: makes session visible to listener thread, # so all attributes above must be initialized first. self._sender.setup_session(self) @@ -1231,8 +2138,16 @@ def disagg_request_id(self) -> int: @property def status(self) -> SessionStatus: + if self._terminal_snapshot is not None: + return self._terminal_snapshot if self._terminal_status is not None: return self._terminal_status + if self._exception is not None or any( + task.status == TaskStatus.ERROR for task in self.kv_tasks + ): + return SessionStatus.ERROR + if self.aux_task is not None and self.aux_task.status == TaskStatus.ERROR: + return SessionStatus.ERROR kv_all_transferred = bool(self.kv_tasks) and all( t.status == TaskStatus.TRANSFERRED for t in self.kv_tasks ) @@ -1259,8 +2174,16 @@ def send(self, slice: KVSlice) -> None: ) task._unique_rid = self.disagg_request_id self.kv_tasks.append(task) + if self._sealed or self._closed: + task.fail( + RuntimeError( + f"TxSession {self.disagg_request_id} is sealed; KV transfer rejected" + ) + ) + return req_info_snapshot = dict(self._sender._get_req_info(task._unique_rid) or {}) - self._sender.dispatch_task(task, req_info_snapshot) + self.begin_dispatch_unlocked() + self._sender.dispatch_task(task, req_info_snapshot, operation_owner=self) def send_aux(self) -> AuxSendTask: with self.lock: @@ -1268,20 +2191,153 @@ def send_aux(self) -> AuxSendTask: task = AuxSendTask(params, self.aux_slot) task._unique_rid = self.disagg_request_id self.aux_task = task + if self._sealed or self._closed: + task.fail( + RuntimeError( + f"TxSession {self.disagg_request_id} is sealed; aux transfer rejected" + ) + ) + return task req_info_snapshot = dict(self._sender._get_req_info(task._unique_rid) or {}) - self._sender.dispatch_task(task, req_info_snapshot) + self.begin_dispatch_unlocked() + self._sender.dispatch_task(task, req_info_snapshot, operation_owner=self) return task - def pack_aux(self, request: LlmRequest) -> None: - """Fill the aux buffer slot with token data from the given request.""" - assert self._aux_buffer is not None, "No aux_buffer set for this session" - assert self.aux_slot is not None, "No aux_slot set for this session" - self._aux_buffer.fill_slot(self.aux_slot, request) + def begin_dispatch_unlocked(self) -> bool: + """Reserve a lifetime credit before preparing peer write metadata.""" + if self._sealed or self._closed: + return False + self._outstanding_operations += 1 + return True - def is_completed(self) -> bool: - """Non-blocking check: has the transfer completed successfully?""" - status = self.status - if self._need_aux: + def register_request_operation( + self, receiver_endpoint: str, request_epoch: Optional[int] + ) -> None: + """Bind a receiver endpoint to one immutable request incarnation.""" + operation = (receiver_endpoint, request_epoch) + with self.lock: + for known_endpoint, known_epoch in self._request_operations: + if ( + known_endpoint == receiver_endpoint + and known_epoch != request_epoch + and (known_epoch is not None or request_epoch is not None) + ): + raise RuntimeError( + "receiver endpoint changed request incarnation inside an active " + f"TxSession: rid={self.disagg_request_id} endpoint={receiver_endpoint} " + f"old_epoch={known_epoch} new_epoch={request_epoch}" + ) + self._request_operations.add(operation) + + def has_request_operation(self, receiver_endpoint: str, request_epoch: Optional[int]) -> bool: + with self.lock: + return (receiver_endpoint, request_epoch) in self._request_operations + + def has_request_endpoint(self, receiver_endpoint: str) -> bool: + with self.lock: + return any( + endpoint == receiver_endpoint + for endpoint, _request_epoch in self._request_operations + ) + + def finish_dispatch(self, write_metas: list[WriteMeta]) -> None: + """Replace one preparation credit with its queued peer obligations.""" + ack_endpoints: list[tuple[str, Optional[int]]] = [] + finalize_close = False + with self.lock: + if self._outstanding_operations <= 0: + raise RuntimeError(f"TxSession {self.disagg_request_id} dispatch credit underflow") + self._outstanding_operations -= 1 + for write_meta in write_metas: + write_meta.operation_owner = self + self._outstanding_operations += len(write_metas) + ack_endpoints, finalize_close = self._post_drain_actions_unlocked() + self._run_post_drain_actions(ack_endpoints, finalize_close, from_worker=False) + + def finish_failed_dispatch(self, notification_count: int) -> None: + """Replace one preparation credit with durable failure notifications. + + Each notification callback retires one credit only after the control + owner has accepted the terminal message. This keeps a strong session + owner across transient ZMQ failures. + """ + ack_endpoints: list[tuple[str, Optional[int]]] = [] + finalize_close = False + with self.lock: + if self._outstanding_operations <= 0: + raise RuntimeError(f"TxSession {self.disagg_request_id} dispatch credit underflow") + self._outstanding_operations -= 1 + self._outstanding_operations += notification_count + ack_endpoints, finalize_close = self._post_drain_actions_unlocked() + self._run_post_drain_actions(ack_endpoints, finalize_close, from_worker=False) + + def retire_operation(self) -> None: + """Retire one peer response obligation after its terminal result was sent.""" + ack_endpoints: list[tuple[str, Optional[int]]] = [] + finalize_close = False + with self.lock: + if self._outstanding_operations <= 0: + raise RuntimeError(f"TxSession {self.disagg_request_id} operation credit underflow") + self._outstanding_operations -= 1 + ack_endpoints, finalize_close = self._post_drain_actions_unlocked() + self._run_post_drain_actions(ack_endpoints, finalize_close, from_worker=True) + + def _post_drain_actions_unlocked( + self, + ) -> tuple[list[tuple[str, Optional[int]]], bool]: + if self._outstanding_operations != 0: + return [], False + ack_operations: list[tuple[str, Optional[int]]] = [] + if self._terminal_status == SessionStatus.CANCELLED: + if self._cancel_ack_operations: + ack_operations = list(self._cancel_ack_operations) + self._cancel_ack_operations.clear() + self._cancel_acked_operations.update(ack_operations) + elif self._cancel_ack_endpoints: + # Legacy drain-ACK compatibility. + ack_operations = [(endpoint, None) for endpoint in self._cancel_ack_endpoints] + self._cancel_ack_endpoints.clear() + self._cancel_acked_endpoints.update(endpoint for endpoint, _epoch in ack_operations) + return ack_operations, self._close_requested and not self._closed + + def _run_post_drain_actions( + self, + ack_operations: list[tuple[str, Optional[int]]], + finalize_close: bool, + *, + from_worker: bool, + ) -> None: + for endpoint, request_epoch in ack_operations: + if not self._sender.send_cancel_ack( + endpoint, + self.disagg_request_id, + request_epoch=request_epoch, + from_worker=from_worker, + ): + # Queue rejection can only happen during teardown. Restore the + # pending obligation so a duplicate cancel can retry rather + # than silently treating an unsent ACK as durable. + with self.lock: + operation = (endpoint, request_epoch) + if request_epoch is None: + self._cancel_acked_endpoints.discard(endpoint) + self._cancel_ack_endpoints.add(endpoint) + else: + self._cancel_acked_operations.discard(operation) + self._cancel_ack_operations.add(operation) + if finalize_close: + self._finalize_close() + + def pack_aux(self, request: LlmRequest) -> None: + """Fill the aux buffer slot with token data from the given request.""" + assert self._aux_buffer is not None, "No aux_buffer set for this session" + assert self.aux_slot is not None, "No aux_slot set for this session" + self._aux_buffer.fill_slot(self.aux_slot, request) + + def is_completed(self) -> bool: + """Non-blocking check: has the transfer completed successfully?""" + status = self.status + if self._need_aux: return status == SessionStatus.FULLY_TRANSFERRED return status in (SessionStatus.KV_TRANSFERRED, SessionStatus.FULLY_TRANSFERRED) @@ -1293,7 +2349,11 @@ def has_failed(self) -> bool: return True return self.aux_task is not None and self.aux_task.status == TaskStatus.ERROR - def cancel(self) -> None: + def cancel( + self, + ack_endpoint: Optional[str] = None, + request_epoch: Optional[int] = None, + ) -> None: """Cancel the session and notify the remote receiver. Safe to call multiple times. TRANSFERRING tasks keep running (mid-write). @@ -1301,25 +2361,118 @@ def cancel(self) -> None: The lock serializes with _deliver_kv_to_agent() so has_transferring_tasks() is accurate the moment this returns. """ + ack_operations: list[tuple[str, Optional[int]]] = [] + notify_receivers = False with self.lock: - if self._terminal_status == SessionStatus.CANCELLED: - return - self._terminal_status = SessionStatus.CANCELLED - exc = RuntimeError(f"TxSession {self.disagg_request_id} cancelled") - for task in self.kv_tasks: - if task.status == TaskStatus.INIT: - task.fail(exc) - if self.aux_task is not None and self.aux_task.status == TaskStatus.INIT: - self.aux_task.fail(exc) + if ack_endpoint is not None: + operation = (ack_endpoint, request_epoch) + if request_epoch is None: + if ack_endpoint not in self._cancel_acked_endpoints: + self._cancel_ack_endpoints.add(ack_endpoint) + elif self._outstanding_operations == 0: + # CANCEL is retransmitted until an ACK arrives. Always + # resend an idempotent ACK after drain; queue acceptance + # is not proof that the previous message was delivered. + ack_operations.append(operation) + elif operation not in self._request_operations: + raise RuntimeError( + "cancel request does not match an active Tx operation: " + f"rid={self.disagg_request_id} endpoint={ack_endpoint} " + f"epoch={request_epoch}" + ) + elif operation not in self._cancel_acked_operations: + self._cancel_ack_operations.add(operation) + elif self._outstanding_operations == 0: + ack_operations.append(operation) + if self._terminal_snapshot is not None: + if self._outstanding_operations == 0: + v2_ack_operations = list(self._cancel_ack_operations) + ack_operations.extend(v2_ack_operations) + self._cancel_ack_operations.clear() + self._cancel_acked_operations.update(v2_ack_operations) + legacy_ack_operations = [ + (endpoint, None) for endpoint in self._cancel_ack_endpoints + ] + ack_operations.extend(legacy_ack_operations) + self._cancel_ack_endpoints.clear() + self._cancel_acked_endpoints.update( + endpoint for endpoint, _epoch in legacy_ack_operations + ) + elif self._terminal_status == SessionStatus.CANCELLED: + newly_ready, _ = self._post_drain_actions_unlocked() + ack_operations.extend(newly_ready) + else: + self._sealed = True + self._terminal_status = SessionStatus.CANCELLED + exc = RuntimeError(f"TxSession {self.disagg_request_id} cancelled") + for task in self.kv_tasks: + if task.status == TaskStatus.INIT: + task.fail(exc) + if self.aux_task is not None and self.aux_task.status == TaskStatus.INIT: + self.aux_task.fail(exc) + newly_ready, _ = self._post_drain_actions_unlocked() + ack_operations.extend(newly_ready) + if self._terminal_snapshot is None and not self._cancel_notified: + self._cancel_notified = True + notify_receivers = True # Send outside the lock to avoid holding it during I/O. - self._sender.send_cancel_to_receivers(self.disagg_request_id) + if notify_receivers: + self._sender.send_cancel_to_receivers(self.disagg_request_id) + self._run_post_drain_actions(ack_operations, False, from_worker=False) def has_transferring_tasks(self) -> bool: - """True if any KV task is currently mid-write (TRANSFERRING). + """True if any KV or aux task is currently mid-write (TRANSFERRING). cancel_request() must return False while this is True. """ - return any(t.status == TaskStatus.TRANSFERRING for t in self.kv_tasks) + with self.lock: + return self._has_transferring_tasks_unlocked() + + def _has_transferring_tasks_unlocked(self) -> bool: + return self._outstanding_operations != 0 + + def seal_and_check_quiescent(self) -> bool: + """Atomically prevent new transfers and check whether active writes drained. + + Sender workers use the same lock for INIT-to-TRANSFERRING transitions, + so no task can start after this method reports quiescence. + """ + with self.lock: + self._sealed = True + return not self._has_transferring_tasks_unlocked() + + def seal_and_snapshot_terminal(self) -> Optional[SessionStatus]: + """Atomically snapshot a drained terminal state without freezing pending work.""" + with self.lock: + if self._terminal_snapshot is not None: + return self._terminal_snapshot + if self._outstanding_operations != 0: + return None + status = self.status + terminal = status in (SessionStatus.ERROR, SessionStatus.CANCELLED) + terminal = terminal or ( + not self._need_aux + and status in (SessionStatus.KV_TRANSFERRED, SessionStatus.FULLY_TRANSFERRED) + ) + terminal = terminal or (self._need_aux and status == SessionStatus.FULLY_TRANSFERRED) + if not terminal: + return None + self._sealed = True + self._terminal_snapshot = status + return status + + def try_mark_transferring(self, task: SendTaskBase) -> bool: + """Start a queued task unless the session has crossed its seal boundary.""" + with self.lock: + if ( + self._sealed + or self._closed + or self._terminal_status in (SessionStatus.ERROR, SessionStatus.CANCELLED) + or task.status not in (TaskStatus.INIT, TaskStatus.TRANSFERRING) + ): + return False + task.status = TaskStatus.TRANSFERRING + return True def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: """Poll or block until KV (and optionally aux) transfer finishes. @@ -1367,12 +2520,17 @@ def set_exception(self, reason: str = ""): if reason: msg += f": {reason}" with self.lock: + if self._terminal_snapshot is not None: + return + if self._terminal_status == SessionStatus.CANCELLED: + return self._exception = RuntimeError(msg) + self._sealed = True self._terminal_status = SessionStatus.ERROR for task in self.kv_tasks: - if not task.is_done: + if task.status == TaskStatus.INIT: task.fail(self._exception) - if self.aux_task is not None and not self.aux_task.is_done: + if self.aux_task is not None and self.aux_task.status == TaskStatus.INIT: self.aux_task.fail(self._exception) @property @@ -1380,15 +2538,27 @@ def exception(self) -> Optional[Exception]: return self._exception def close(self): - if getattr(self, "_closed", False): - return - self._closed = True + finalize = False + with self.lock: + if getattr(self, "_closed", False): + return + self._sealed = True + self._close_requested = True + finalize = self._outstanding_operations == 0 + if finalize: + self._finalize_close() + + def _finalize_close(self) -> None: + with self.lock: + if self._closed or self._outstanding_operations != 0: + return + self._closed = True if self._aux_buffer is not None and self.aux_slot is not None: self._aux_buffer.free_slot(self.aux_slot) self.aux_slot = None # Unregister from Sender; keep fields alive for in-flight worker threads. if self._sender is not None: - self._sender.clear_session(self.disagg_request_id) + self._sender.clear_session(self.disagg_request_id, set(self._request_operations)) def __enter__(self): return self @@ -1456,23 +2626,39 @@ def print_perf_info(self, peer_rank: int, instance_name: str, instance_rank: int class Receiver(ReceiverBase): + _TOMBSTONE_LIMIT = 65_536 + _SHUTDOWN_DRAIN_LOG_INTERVAL_S = 5.0 + def __init__( self, peer_registrar: PeerRegistrar, agent: BaseTransferAgent, bounce=None, ): + self._shutdown = False + self._shutdown_lock = threading.Lock() + self._shutdown_complete = threading.Event() + self._shutdown_thread: Optional[threading.Thread] = None + self._shutdown_error: Optional[BaseException] = None self._registrar = peer_registrar self._agent = agent self._bounce = bounce - self._dealers = {} + self._control = _ControlPlane("native-receiver") self._sender_ep_instance_map = {} + self._sender_info_capabilities: dict[str, _NativeProtocolCapabilities] = {} + self._sender_endpoint_capabilities: dict[str, _NativeProtocolCapabilities] = {} self._messenger = ZMQMessenger(mode="ROUTER") self._sessions = {} # unique_rid -> RxSession - self._sessions_lock = threading.Lock() # Protects _sessions and _pre_cancelled_rids - self._pre_cancelled_rids: set[int] = set() - self._shutdown = False + self._sessions_lock = threading.Lock() + self._sessions_drained = threading.Condition(self._sessions_lock) + self._draining_sessions: dict[int, RxSession] = {} + # A pre-cancel is live protocol state: evicting it can let a delayed + # request start after cancellation was already acknowledged upstream. + self._pre_cancelled_rids: dict[int, None] = {} + self._closed_rids: OrderedDict[int, None] = OrderedDict() + self._closed_operations: OrderedDict[tuple[int, int], None] = OrderedDict() + self._protocol_error: Optional[RuntimeError] = None self._start_listener() logger.info(f"Receiver init with endpoint: {self._messenger.endpoint}") @@ -1481,32 +2667,140 @@ def __init__( def endpoint(self): return self._messenger.endpoint - def shutdown(self): - if getattr(self, "_shutdown", False): - return - self._shutdown = True - for dealer in self._dealers.values(): - try: - dealer.stop() - except Exception as e: - logger.warning(f"Failed to stop dealer during Receiver shutdown: {e}") - self._dealers.clear() - self._messenger.stop() + def shutdown(self) -> Optional[threading.Event]: + if self._is_progress_thread(): + return self._defer_shutdown() + run_shutdown = False + with self._shutdown_lock: + if self._shutdown_complete.is_set(): + self._raise_shutdown_error() + return + if not self._shutdown: + self._shutdown = True + run_shutdown = True - def clear_session(self, unique_rid: int): - with self._sessions_lock: + if run_shutdown: + self._run_shutdown() + else: + self._shutdown_complete.wait() + self._raise_shutdown_error() + return None + + def _is_progress_thread(self) -> bool: + """Whether blocking here would prevent an Rx operation from draining.""" + current = threading.current_thread() + listener_thread = getattr(self._messenger, "_listener_thread", None) + scatter_thread = getattr(self._bounce, "_scatter_thread", None) + return current is listener_thread or current is scatter_thread + + def _defer_shutdown(self) -> threading.Event: + """Let the current listener/scatter callback retire its own drain credit.""" + with self._shutdown_lock: + if self._shutdown_complete.is_set() or self._shutdown: + return self._shutdown_complete + self._shutdown = True + self._shutdown_thread = threading.Thread( + target=self._run_shutdown, + name="trtllm-rx-shutdown", + daemon=True, + ) + shutdown_thread = self._shutdown_thread + logger.warning( + "Receiver shutdown was requested from an Rx progress thread; " + "deferring teardown until the callback returns and all sessions drain" + ) + shutdown_thread.start() + return self._shutdown_complete + + def _run_shutdown(self) -> None: + try: + self._close_and_drain_sessions() + # The listener and control dealers are themselves part of the + # receive progress engine. Stop them only after no session can + # still need a terminal result or cancel ACK. + self._messenger.stop() + self._control.flush() + self._control.shutdown() + except BaseException as error: + self._shutdown_error = error + logger.exception("Receiver shutdown failed before receive resources were quiescent") + finally: + self._shutdown_complete.set() + + def _close_and_drain_sessions(self) -> None: + # Closing every still-live session makes it either finalize immediately + # or retain itself strongly in _draining_sessions. Snapshot first: a + # finalizer calls clear_session(), which takes the same condition lock. + with self._sessions_drained: + sessions = [] + for session_ref in self._sessions.values(): + session = session_ref() + if session is not None: + sessions.append(session) + for session in sessions: + session.close() + + with self._sessions_drained: + while self._draining_sessions: + draining_rids = list(self._draining_sessions) + logger.warning( + "Receiver shutdown is waiting for %d session(s) to drain before " + "stopping its listener or releasing transfer memory: %s", + len(draining_rids), + draining_rids, + ) + self._sessions_drained.wait(timeout=self._SHUTDOWN_DRAIN_LOG_INTERVAL_S) + + def _raise_shutdown_error(self) -> None: + if self._shutdown_error is not None: + raise RuntimeError( + "Receiver shutdown did not complete safely" + ) from self._shutdown_error + + def clear_session(self, unique_rid: int, request_epoch: Optional[int] = None): + with self._sessions_drained: self._sessions.pop(unique_rid, None) + self._draining_sessions.pop(unique_rid, None) + self._remember_tombstone(self._closed_rids, unique_rid) + if request_epoch is not None: + self._remember_tombstone(self._closed_operations, (unique_rid, request_epoch)) + self._sessions_drained.notify_all() + + def retain_draining_session(self, session: "RxSession") -> None: + with self._sessions_drained: + self._draining_sessions[session.disagg_request_id] = session + self._sessions_drained.notify_all() def setup_session(self, rx_session: RxSessionBase): pre_cancel = False with self._sessions_lock: + if self._shutdown: + raise RuntimeError("Cannot create an RxSession after Receiver shutdown started") + if self._protocol_error is not None: + raise RuntimeError( + "Receiver is in fail-stop protocol state" + ) from self._protocol_error self._sessions[rx_session.disagg_request_id] = weakref.ref(rx_session) if rx_session.disagg_request_id in self._pre_cancelled_rids: pre_cancel = True - self._pre_cancelled_rids.discard(rx_session.disagg_request_id) + self._pre_cancelled_rids.pop(rx_session.disagg_request_id, None) + self._closed_rids.pop(rx_session.disagg_request_id, None) if pre_cancel: rx_session.cancel() + @classmethod + def _remember_tombstone(cls, tombstones: OrderedDict, key) -> None: + tombstones.pop(key, None) + tombstones[key] = None + while len(tombstones) > cls._TOMBSTONE_LIMIT: + tombstones.popitem(last=False) + + def _fail_protocol(self, error: RuntimeError) -> None: + with self._sessions_lock: + if self._protocol_error is None: + self._protocol_error = error + logger.error("Native receiver entered fail-stop protocol state: %s", error) + def _get_session(self, unique_rid: Optional[int]) -> Optional["RxSession"]: with self._sessions_lock: session_ref = self._sessions.get(unique_rid) @@ -1547,6 +2841,54 @@ def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo: slice_id=task.slice_id, ) + def _build_bounce_destination_plan( + self, receiver_req: RecvReqInfo, peer_ri: RankInfo + ) -> tuple[np.ndarray, np.ndarray]: + """Derive this sender rank's exact local destination sequence. + + The receiver-side registrar builds the reverse KV mapping (local GEN + cache to the peer CTX rank). Its source fragments are therefore the + exact local byte ranges that the same peer writes in the forward + direction. Dummy peer slots provide only matching cardinality; their + addresses do not affect the returned local source fragments. + """ + extractor = self._registrar.self_extractor + peer_extractor = self._registrar.peer_extractor( + peer_ri.instance_name, peer_ri.instance_rank + ) + dst_parts: list[np.ndarray] = [] + size_specs: list[tuple[int, int]] = [] + for (self_lg, self_pi), (peer_lg, peer_pi) in self._registrar.get_pool_mapping( + peer_ri + ).items(): + block_ids = np.asarray(receiver_req.block_ids_per_layer_groups[self_lg], dtype=np.int64) + block_ids = block_ids[block_ids >= 0] + if block_ids.size == 0: + continue + self_region = extractor.extract(block_ids, layer_group_id=self_lg, pool_idx=self_pi) + peer_region = peer_extractor.extract( + np.zeros(block_ids.size, dtype=np.int64), + layer_group_id=peer_lg, + pool_idx=peer_pi, + ) + mapper = self._registrar.get_kv_map(peer_ri, (self_lg, self_pi), (peer_lg, peer_pi)) + region_pair = mapper.map(self_region, peer_region) + region_pairs = region_pair if isinstance(region_pair, list) else [region_pair] + for pair in region_pairs: + dst_parts.append(pair.src.memory.ptrs) + size_specs.append((pair.src.memory.ptrs.size, pair.src.memory.bytes_per_region)) + + if not dst_parts: + return np.array([], dtype=np.int64), np.array([], dtype=np.int64) + dst_ptrs = np.concatenate(dst_parts) + counts, values = zip(*size_specs) + sizes = np.repeat(np.array(values, dtype=np.int64), counts) + order = sorted(range(dst_ptrs.size), key=lambda index: int(dst_ptrs[index])) + return ( + np.asarray([dst_ptrs[index] for index in order], dtype=np.int64), + np.asarray([sizes[index] for index in order], dtype=np.int64), + ) + @staticmethod def _fanin_bounce_safe(overlap, peer_ri) -> bool: """Whether multi-writer bounce's equal total//num_writers split is valid for this overlap. @@ -1585,6 +2927,12 @@ def dispatch_task(self, task: KVRecvTask): sender_dp_rank = params.ctx_dp_rank peer_infos: RankInfo = self._get_sender_info(params) + if sender_dp_rank is None and _requires_drain_ack_protocol(): + raise RuntimeError( + "asynchronous native-transfer consensus does not support ADP " + "broadcast because no single writer cohort is known before dispatch" + ) + if sender_dp_rank is not None: # Normal path: ctx_dp_rank is known, send to overlapping ranks. peer_overlap = self._registrar.get_peer_overlap(peer_infos, sender_dp_rank) @@ -1620,29 +2968,144 @@ def dispatch_task(self, task: KVRecvTask): allow_bounce = task.expected_transfers == 1 or ( sender_dp_rank is not None and self._fanin_bounce_safe(topo_overlap, peer_infos) ) - bounced = allow_bounce and self._bounce.reserve(receiver_req, task.expected_transfers) + expected_destination_plans = None + if allow_bounce and self._bounce.enabled: + try: + expected_destination_plans = { + rank: self._build_bounce_destination_plan( + receiver_req, + self._registrar.get_peer_rank_info(peer_infos.instance_name, rank), + ) + for rank in peer_overlap.ranks + } + except (AssertionError, IndexError, KeyError, TypeError, ValueError) as error: + logger.warning( + "KV bounce cannot derive an exact receiver-owned destination plan; " + "using per-fragment transfer: %s", + error, + ) + allow_bounce = False + bounced = allow_bounce and self._bounce.reserve( + receiver_req, + task.expected_transfers, + expected_destination_plans=expected_destination_plans, + ) session = self._get_session(task._unique_rid) if session is None: + self._bounce.release_idle_reservation((receiver_req.unique_rid, receiver_req.slice_id)) raise RuntimeError( f"dispatch_task: RxSession {task._unique_rid} not found; " "session may have been closed before dispatch" ) - session.mark_transferring(task.slice_id) - # Cache sender endpoints so cancel() can send CANCEL_SESSION to them. - session._sender_endpoints.update( - peer_infos.sender_endpoints[rank] for rank in peer_overlap.ranks + endpoint_by_rank = {rank: peer_infos.sender_endpoints[rank] for rank in peer_overlap.ranks} + sender_endpoints = set(endpoint_by_rank.values()) + # Do not mix ACK-authoritative and legacy result-authoritative cleanup + # inside one fan-out. A rolling deployment therefore stays entirely + # on the legacy path until every participating endpoint advertises + # drain ACK support. + ack_capable_endpoints = ( + sender_endpoints + if sender_endpoints + and all( + _supports_native_protocol_v2( + self._sender_endpoint_capabilities.get(endpoint, _LEGACY_PROTOCOL_CAPABILITIES) + ) + for endpoint in sender_endpoints + ) + else set() ) + v2_enabled = bool(ack_capable_endpoints) and sender_dp_rank is not None + receiver_req.request_epoch = session.request_epoch if v2_enabled else None + allowed_rank_cohorts = None + if sender_dp_rank is None: + allowed_rank_cohorts = tuple( + frozenset(self._registrar.get_peer_overlap(peer_infos, dp).ranks) + for dp in range(peer_infos.dp_size) + ) + if not session.mark_transferring( + task.slice_id, + endpoint_by_rank, + ack_capable_endpoints, + request_epoch=receiver_req.request_epoch, + bounced=bounced, + allowed_rank_cohorts=allowed_rank_cohorts, + ): + # A terminal path sealed the session after receive() reserved the + # task but before any REQUEST_DATA message was sent. + self._bounce.release_idle_reservation((receiver_req.unique_rid, receiver_req.slice_id)) + if not task.is_done: + task.fail( + RuntimeError( + f"RxSession {receiver_req.unique_rid} sealed before transfer started" + ) + ) + return # Fan-in: each sender gets its own sub-region base (writers must not overwrite); else serialize once. fanin_bounce = bounced and task.expected_transfers > 1 key = (receiver_req.unique_rid, receiver_req.slice_id) - receiver_req_bytes = None if fanin_bounce else receiver_req.to_bytes() - for i, rank in enumerate(peer_overlap.ranks): - if task._perf_timer is not None: - task._perf_timer.record_task_start(rank) - if fanin_bounce: - receiver_req.bounce_dst_base = self._bounce.writer_base(key, i) + receiver_req_bytes = None + sent_ranks: set[int] = set() + try: + if bounced: + # Install the lifetime callback before advertising any address. + # Result callbacks are too late for cancel-before-result and + # malformed-result paths; both must still retire this slice's + # settlement credit after exact drain proof. + self._bounce.set_completion_callback( + key, session._make_bounce_settlement_callback(task) + ) + if not fanin_bounce: + if bounced: + only_rank = peer_overlap.ranks[0] + receiver_req.bounce_dst_base = self._bounce.bind_writer(key, only_rank, 0) receiver_req_bytes = receiver_req.to_bytes() - self._request_sender_data(peer_infos.sender_endpoints[rank], receiver_req_bytes) + for i, rank in enumerate(peer_overlap.ranks): + if task._perf_timer is not None: + task._perf_timer.record_task_start(rank) + if fanin_bounce: + receiver_req.bounce_dst_base = self._bounce.bind_writer(key, rank, i) + receiver_req_bytes = receiver_req.to_bytes() + endpoint = endpoint_by_rank[rank] + assert receiver_req_bytes is not None + self._request_sender_data(endpoint, receiver_req_bytes) + sent_ranks.add(rank) + except Exception as error: + logger.error( + "REQUEST_DATA fan-out failed for request %s slice=%s after %s/%s endpoint(s): %s", + receiver_req.unique_rid, + receiver_req.slice_id, + len(sent_ranks), + len(sender_endpoints), + error, + ) + ( + cancel_endpoints, + cancel_operations, + orphaned_bounces, + ) = session.fail_partial_dispatch( + task.slice_id, + sent_ranks, + error, + ) + for orphaned_key in orphaned_bounces: + try: + if orphaned_key == key and not sent_ranks: + self._bounce.release_idle_reservation(orphaned_key) + session._retire_bounce_settlement(task.slice_id) + else: + self._bounce.orphan_reservation(orphaned_key) + except Exception as bounce_error: + logger.error( + "Failed to retain partial-dispatch bounce reservation %s: %s", + orphaned_key, + bounce_error, + ) + if cancel_endpoints: + self.send_cancel_to_senders( + receiver_req.unique_rid, + cancel_endpoints, + cancel_operations, + ) return @staticmethod @@ -1656,12 +3119,24 @@ def _should_register_peer(self, params: DisaggregatedParams) -> bool: endpoint = self._extract_info_endpoint(params) return endpoint not in self._sender_ep_instance_map - def _get_or_connect_dealer(self, endpoint: Optional[str]): - if endpoint is None: - raise ValueError("Receiver: peer endpoint is None; peer may not have registered yet") - if endpoint not in self._dealers: - self._dealers[endpoint] = ZMQMessenger(mode="DEALER", endpoint=endpoint) - return self._dealers[endpoint] + def _send_control_message( + self, + endpoint: Optional[str], + message: list[bytes], + *, + retry: bool = False, + wait: bool = True, + on_sent: Optional[Callable[[], None]] = None, + repeat_until: Optional[Callable[[], bool]] = None, + ) -> _ControlSend: + return self._control.send( + endpoint, + message, + retry=retry, + wait=wait, + on_sent=on_sent, + repeat_until=repeat_until, + ) def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: info_endpoint = self._extract_info_endpoint(params) @@ -1672,29 +3147,95 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: messenger.send([MessageType.REQUEST_INSTANCE_INFO]) message = messenger.receive() sender_info = RankInfo.from_bytes(message[0]) + capabilities = _decode_protocol_capabilities( + message[1] if len(message) > 1 else None + ) finally: messenger.stop() + if _requires_drain_ack_protocol() and not _supports_native_protocol_v2(capabilities): + raise RuntimeError( + "asynchronous Python-transceiver consensus requires native-transfer " + f"protocol version {_NATIVE_PROTOCOL_VERSION}, but the context peer at " + f"{info_endpoint} advertised legacy protocol version {capabilities.version}" + ) + for endpoint in sender_info.sender_endpoints: - dealer = self._get_or_connect_dealer(endpoint) rank_info = self._registrar.self_rank_info - dealer.send([MessageType.REGISTER_RANK_INFO, rank_info.to_bytes()]) + self._send_control_message( + endpoint, + [ + MessageType.REGISTER_RANK_INFO, + rank_info.to_bytes(), + _encode_protocol_capabilities(), + ], + ) + self._sender_endpoint_capabilities[endpoint] = capabilities self._sender_ep_instance_map[info_endpoint] = sender_info + self._sender_info_capabilities[info_endpoint] = capabilities return sender_info else: return self._sender_ep_instance_map[info_endpoint] - def send_cancel_to_senders(self, unique_rid: int, sender_endpoints: set[str]) -> None: - """Notify all senders involved in this session to cancel.""" + def send_cancel_to_senders( + self, + unique_rid: int, + sender_endpoints: set[str], + cancel_operations: set[tuple[str, int]], + ) -> None: + """Notify senders and retransmit v2 CANCEL until its exact ACK arrives.""" + errors: list[Exception] = [] for endpoint in sender_endpoints: try: - self._get_or_connect_dealer(endpoint).send( - [MessageType.CANCEL_SESSION, str(unique_rid).encode("ascii")] - ) + message = [MessageType.CANCEL_SESSION, str(unique_rid).encode("ascii")] + epochs = { + epoch + for operation_endpoint, epoch in cancel_operations + if operation_endpoint == endpoint + } + if len(epochs) > 1: + raise RuntimeError( + f"multiple active request epochs for rid={unique_rid} endpoint={endpoint}" + ) + if epochs: + request_epoch = next(iter(epochs)) + message.extend( + [ + self.endpoint.encode("utf-8"), + str(request_epoch).encode("ascii"), + ] + ) + + def acked( + endpoint=endpoint, + request_epoch=request_epoch, + unique_rid=unique_rid, + ) -> bool: + session = self._get_session(unique_rid) + return session is None or not session.is_cancel_pending( + endpoint, request_epoch + ) + + self._send_control_message( + endpoint, + message, + retry=True, + wait=False, + repeat_until=acked, + ) + else: + self._send_control_message(endpoint, message, retry=True, wait=False) except Exception as e: - logger.warning(f"send_cancel_to_senders: failed for rid={unique_rid}: {e}") + errors.append(e) + logger.error(f"send_cancel_to_senders: failed for rid={unique_rid}: {e}") + if errors: + error = RuntimeError( + f"failed to enqueue cancellation for request {unique_rid}; retaining transfer state" + ) + self._fail_protocol(error) + raise error from errors[0] def _start_listener(self): def handle_message(messages: list[bytes]) -> bool: @@ -1718,6 +3259,11 @@ def handle_message(messages: list[bytes]) -> bool: self._handle_cancel_session(msg) except Exception as e: logger.error(f"Receiver: error handling CANCEL_SESSION: {e}") + case MessageType.CANCEL_SESSION_ACK: + try: + self._handle_cancel_session_ack(msg) + except Exception as e: + logger.error(f"Receiver: error handling CANCEL_SESSION_ACK: {e}") case _: logger.error(f"Receiver received unknown message type: {msg[0]}") return True @@ -1726,32 +3272,101 @@ def handle_message(messages: list[bytes]) -> bool: def _handle_cancel_session(self, message: list[bytes]): unique_rid = int(message[1]) + sender_endpoint = message[2].decode("utf-8") if len(message) > 2 else None + request_epoch = int(message[3]) if len(message) > 3 else None session = None with self._sessions_lock: session_ref = self._sessions.get(unique_rid) if session_ref is None: - self._pre_cancelled_rids.add(unique_rid) + if unique_rid not in self._closed_rids: + if len(self._pre_cancelled_rids) >= _LIVE_PROTOCOL_STATE_LIMIT: + error = RuntimeError( + "Receiver pre-cancel state reached its safety limit; " + "refusing to evict live state" + ) + self._protocol_error = error + raise error + self._pre_cancelled_rids[unique_rid] = None else: session = session_ref() if session is None: - self._pre_cancelled_rids.add(unique_rid) + if unique_rid not in self._closed_rids: + if len(self._pre_cancelled_rids) >= _LIVE_PROTOCOL_STATE_LIMIT: + error = RuntimeError( + "Receiver pre-cancel state reached its safety limit; " + "refusing to evict live state" + ) + self._protocol_error = error + raise error + self._pre_cancelled_rids[unique_rid] = None if session is not None: + if request_epoch is not None: + session.validate_remote_cancel(sender_endpoint, request_epoch) session.cancel() + def _handle_cancel_session_ack(self, message: list[bytes]) -> None: + unique_rid = int(message[1]) + sender_endpoint = message[2].decode("utf-8") + request_epoch = int(message[3]) if len(message) > 3 else None + session = self._get_session(unique_rid) + if session is None: + if request_epoch is not None: + with self._sessions_lock: + if (unique_rid, request_epoch) in self._closed_operations: + return + self._fail_protocol( + RuntimeError( + f"cancel ACK for unknown request incarnation rid={unique_rid} " + f"epoch={request_epoch}" + ) + ) + return + session.process_cancel_ack(sender_endpoint, request_epoch) + def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): if message[0] != MessageType.KV_AGENT_RESULT: logger.error( f"_process_kv_agent_result: unexpected msg_type={message[0]!r}, expected KV_AGENT_RESULT" ) return - peer_rank, unique_rid, sender_slice_id, is_last_slice, status_code, transfer_size = ( - _KV_RESULT_PREFIX.unpack(message[1]) - ) + sender_endpoint = None + request_epoch = None + tail_index = 2 + if len(message[1]) == _KV_RESULT_PREFIX_V2.size: + ( + peer_rank, + unique_rid, + request_epoch, + sender_slice_id, + is_last_slice, + status_code, + transfer_size, + ) = _KV_RESULT_PREFIX_V2.unpack(message[1]) + if len(message) < 3: + raise RuntimeError("v2 KV result is missing sender endpoint") + sender_endpoint = message[2].decode("utf-8") + tail_index = 3 + elif len(message[1]) == _KV_RESULT_PREFIX.size: + peer_rank, unique_rid, sender_slice_id, is_last_slice, status_code, transfer_size = ( + _KV_RESULT_PREFIX.unpack(message[1]) + ) + else: + raise RuntimeError(f"invalid KV result prefix size {len(message[1])}") from .bounce import decode_result_tail - dst_ptrs, sizes, src_base = decode_result_tail(message) + dst_ptrs, sizes, src_base = decode_result_tail(message, tail_index=tail_index) session = self._get_session(unique_rid) if session is None: + if request_epoch is not None: + with self._sessions_lock: + if (unique_rid, request_epoch) in self._closed_operations: + return + self._fail_protocol( + RuntimeError( + f"KV result for unknown request incarnation rid={unique_rid} " + f"epoch={request_epoch}" + ) + ) logger.warning( f"_process_kv_agent_result: session {unique_rid} not found (already closed?), dropping status" ) @@ -1765,25 +3380,45 @@ def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): sizes=sizes, src_base=src_base, transfer_size=transfer_size, + sender_endpoint=sender_endpoint, + request_epoch=request_epoch, ) def _process_aux_agent_result(self, _send_id: bytes, message: list[bytes]): - _msg_type, peer_rank, unique_rid, status = decode_message(message) - peer_rank = int(peer_rank) - unique_rid = int(unique_rid) + if len(message) not in (4, 6): + raise RuntimeError(f"invalid AUX result frame count {len(message)}") + peer_rank = int(message[1]) + unique_rid = int(message[2]) + status = message[3].decode("utf-8") + sender_endpoint = message[4].decode("utf-8") if len(message) == 6 else None + request_epoch = int(message[5]) if len(message) == 6 else None session = self._get_session(unique_rid) if session is None: + if request_epoch is not None: + with self._sessions_lock: + if (unique_rid, request_epoch) in self._closed_operations: + return + self._fail_protocol( + RuntimeError( + f"AUX result for unknown request incarnation rid={unique_rid} " + f"epoch={request_epoch}" + ) + ) logger.warning( f"_process_aux_agent_result: session {unique_rid} not found (already closed?), dropping status" ) return - session.process_aux_agent_result(peer_rank, AgentResult(status)) + session.process_aux_agent_result( + peer_rank, + AgentResult(status), + sender_endpoint=sender_endpoint, + request_epoch=request_epoch, + ) def _request_sender_data(self, endpoint: str, receiver_info_bytes: bytes): # receiver_info serialized once and reused for every peer rank (block-table msgpack isn't free at fan-out). logger.debug("Sending data request to endpoint '%s'", endpoint) - messenger = self._get_or_connect_dealer(endpoint) - messenger.send([MessageType.REQUEST_DATA, receiver_info_bytes]) + self._send_control_message(endpoint, [MessageType.REQUEST_DATA, receiver_info_bytes]) def __del__(self): try: @@ -1829,7 +3464,23 @@ def __init__( self._aux_count = 0 self._aux_status: TaskStatus = TaskStatus.INIT self._sender_endpoints: set[str] = set() + self.request_epoch = secrets.randbits(63) or 1 self.lock = threading.Lock() + self._sealed = False + self._outstanding_operations = 0 + self._expected_operations: dict[tuple, _ExpectedReceiveOperation] = {} + self._retired_operation_keys: set[tuple] = set() + self._pending_scatter_callbacks = 0 + self._pending_bounce_settlements: set[int] = set() + self._aux_obligations_reserved = False + self._aux_obligation_slice_id: Optional[int] = None + self._cancel_pending_endpoints: set[str] = set() + self._cancel_pending_operations: set[tuple[str, int]] = set() + self._ack_capable_endpoints: set[str] = set() + self._v2_enabled: Optional[bool] = None + self._legacy_rank_cohorts: dict[int, tuple[frozenset[int], ...]] = {} + self._legacy_bound_cohorts: dict[int, frozenset[int]] = {} + self._close_requested = False self._receiver.setup_session(self) @property @@ -1860,23 +3511,278 @@ def status(self) -> SessionStatus: return SessionStatus.TRANSFERRING return SessionStatus.INIT - def mark_transferring(self, slice_id: int): + def mark_transferring( + self, + slice_id: int, + endpoint_by_rank: dict[int, str] | set[str], + ack_capable_endpoints: Optional[set[str]] = None, + *, + request_epoch: Optional[int] = None, + bounced: bool = False, + allowed_rank_cohorts: Optional[tuple[frozenset[int], ...]] = None, + ) -> bool: + """Start a receive task and bind every advertised remote operation.""" + with self.lock: + task = self._kv_tasks[slice_id] + if ( + self._sealed + or self._closed + or self._terminal_status in (SessionStatus.ERROR, SessionStatus.CANCELLED) + or task.status != TaskStatus.INIT + ): + return False + if isinstance(endpoint_by_rank, set): + # Compatibility for legacy callers/tests that predate rank + # identity binding. Production dispatch always passes a map. + sender_endpoints = set(endpoint_by_rank) + rank_endpoints: dict[int, str] = {} + else: + rank_endpoints = dict(endpoint_by_rank) + sender_endpoints = set(rank_endpoints.values()) + v2_enabled = request_epoch is not None + if self._v2_enabled is None: + self._v2_enabled = v2_enabled + elif self._v2_enabled != v2_enabled: + raise RuntimeError( + f"RxSession {self.disagg_request_id} cannot mix native protocol versions" + ) + prior_all_ack_capable = ( + not self._sender_endpoints or self._ack_capable_endpoints == self._sender_endpoints + ) + this_all_ack_capable = bool(sender_endpoints) and ( + ack_capable_endpoints == sender_endpoints + ) + self._sender_endpoints.update(sender_endpoints) + if prior_all_ack_capable and this_all_ack_capable: + self._ack_capable_endpoints.update(sender_endpoints) + else: + # A mixed-version session must use legacy terminal results for + # every peer; otherwise one ACK could retire another peer's + # still-active write. + self._ack_capable_endpoints.clear() + task.status = TaskStatus.TRANSFERRING + if allowed_rank_cohorts is not None: + self._legacy_rank_cohorts[slice_id] = allowed_rank_cohorts + # The concrete ADP cohort is selected by the first terminal + # result. Credits remain count-based until that immutable bind. + self._outstanding_operations += task.expected_transfers + else: + for rank, endpoint in rank_endpoints.items(): + operation_key = ("kv", slice_id, rank) + if operation_key in self._expected_operations: + raise RuntimeError( + f"duplicate receive operation {operation_key} for " + f"request {self.disagg_request_id}" + ) + self._expected_operations[operation_key] = _ExpectedReceiveOperation( + endpoint, request_epoch + ) + self._outstanding_operations += ( + len(rank_endpoints) if rank_endpoints else task.expected_transfers + ) + if self._need_aux and not self._aux_obligations_reserved: + if allowed_rank_cohorts is None: + for rank, endpoint in rank_endpoints.items(): + self._expected_operations[("aux", rank)] = _ExpectedReceiveOperation( + endpoint, request_epoch + ) + self._outstanding_operations += ( + len(rank_endpoints) if rank_endpoints else task.expected_transfers + ) + else: + self._outstanding_operations += task.expected_transfers + self._aux_obligations_reserved = True + self._aux_obligation_slice_id = slice_id + if bounced: + self._pending_bounce_settlements.add(slice_id) + self._pending_scatter_callbacks = len(self._pending_bounce_settlements) + return True + + def _bind_legacy_cohort_unlocked(self, slice_id: int, peer_rank: int) -> None: + cohorts = self._legacy_rank_cohorts.get(slice_id) + if cohorts is None: + return + bound = self._legacy_bound_cohorts.get(slice_id) + if bound is None: + matches = [cohort for cohort in cohorts if peer_rank in cohort] + if len(matches) != 1: + raise RuntimeError( + f"rank {peer_rank} does not identify one allowed ADP cohort " + f"for request {self.disagg_request_id} slice={slice_id}" + ) + bound = matches[0] + other_bound_cohorts = set(self._legacy_bound_cohorts.values()) + if other_bound_cohorts and other_bound_cohorts != {bound}: + raise RuntimeError( + f"rank {peer_rank} selects a different ADP writer cohort " + f"for request {self.disagg_request_id} slice={slice_id}" + ) + self._legacy_bound_cohorts[slice_id] = bound + if peer_rank not in bound: + raise RuntimeError( + f"rank {peer_rank} is outside the bound ADP writer cohort " + f"for request {self.disagg_request_id} slice={slice_id}" + ) + + def _validate_receive_operation_unlocked( + self, + operation_key: tuple, + *, + sender_endpoint: Optional[str], + request_epoch: Optional[int], + ) -> str: + """Return ``accept``/``duplicate`` or raise without retiring credit.""" + if operation_key in self._retired_operation_keys: + return "duplicate" + kind = operation_key[0] + if kind == "kv" and operation_key[1] in self._legacy_rank_cohorts: + self._bind_legacy_cohort_unlocked(operation_key[1], operation_key[2]) + return "accept" + if kind == "aux" and self._legacy_rank_cohorts: + # Aux may arrive before KV. Bind it through the same immutable + # cohort selection instead of accepting it count-only and letting + # a later KV result select a different DP group. + if self._aux_obligation_slice_id is None: + raise RuntimeError( + f"aux result has no owning KV slice for request {self.disagg_request_id}" + ) + self._bind_legacy_cohort_unlocked(self._aux_obligation_slice_id, operation_key[1]) + return "accept" + expected = self._expected_operations.get(operation_key) + if expected is None: + # Only compatibility/unit-test sessions created without production + # rank maps may use count-based legacy retirement. + if not self._expected_operations and request_epoch is None: + return "accept" + raise RuntimeError( + f"unexpected native-transfer result operation {operation_key} " + f"for request {self.disagg_request_id}" + ) + if expected.request_epoch != request_epoch: + raise RuntimeError( + f"native-transfer result epoch mismatch for {operation_key}: " + f"expected {expected.request_epoch}, got {request_epoch}" + ) + if request_epoch is not None and expected.sender_endpoint != sender_endpoint: + raise RuntimeError( + f"native-transfer result source mismatch for {operation_key}: " + f"expected {expected.sender_endpoint}, got {sender_endpoint}" + ) + return "accept" + + def _fail_protocol_unlocked(self, error: RuntimeError) -> set[tuple[str, int]]: + self._sealed = True + self._exception = error + if self._terminal_status != SessionStatus.CANCELLED: + self._terminal_status = SessionStatus.ERROR + targets = { + (expected.sender_endpoint, expected.request_epoch) + for key, expected in self._expected_operations.items() + if key not in self._retired_operation_keys and expected.request_epoch is not None + } + self._cancel_pending_operations.update(targets) + self._cancel_pending_endpoints.update(endpoint for endpoint, _epoch in targets) + return targets + + def fail_partial_dispatch( + self, + slice_id: int, + sent_ranks: set[int], + error: Exception, + ) -> tuple[set[str], set[tuple[str, int]], list[tuple[int, int]]]: + """Seal a partially published fan-out without inventing remote quiescence. + + Endpoints which never received REQUEST_DATA cannot write, so their + reserved receive credits are retired locally. Endpoints which did + receive it remain protected by either the negotiated drain ACK or + their legacy terminal result. In the ambiguous legacy ADP-broadcast + case, cleanup deliberately remains retained rather than guessing. + """ with self.lock: - self._kv_tasks[slice_id].status = TaskStatus.TRANSFERRING + task = self._kv_tasks[slice_id] + self._sealed = True + if self._terminal_status != SessionStatus.CANCELLED: + self._terminal_status = SessionStatus.ERROR + self._exception = RuntimeError( + f"REQUEST_DATA fan-out failed for request {self.disagg_request_id}: {error}" + ) + task.fail(self._exception) + + if self._expected_operations: + never_dispatched_keys = { + key + for key in self._expected_operations + if (key[0] == "kv" and key[1] == slice_id and key[2] not in sent_ranks) + or ( + key[0] == "aux" + and self._aux_obligation_slice_id == slice_id + and key[1] not in sent_ranks + ) + } + for operation_key in never_dispatched_keys: + self._retire_receive_operation_unlocked(operation_key) + sent_operations = { + expected + for key, expected in self._expected_operations.items() + if key not in self._retired_operation_keys + and ( + (key[0] == "kv" and key[1] == slice_id and key[2] in sent_ranks) + or ( + key[0] == "aux" + and self._aux_obligation_slice_id == slice_id + and key[1] in sent_ranks + ) + ) + } + cancel_endpoints = {operation.sender_endpoint for operation in sent_operations} + cancel_operations = { + (operation.sender_endpoint, operation.request_epoch) + for operation in sent_operations + if operation.request_epoch is not None + } + else: + max_remote_writers = min(task.expected_transfers, len(sent_ranks)) + never_dispatched = task.expected_transfers - max_remote_writers + retired = never_dispatched + if self._need_aux and self._aux_obligation_slice_id == slice_id: + retired += never_dispatched + if retired > self._outstanding_operations: + raise RuntimeError( + f"RxSession {self.disagg_request_id} partial-dispatch credit underflow" + ) + self._outstanding_operations -= retired + cancel_endpoints = set(self._sender_endpoints) + cancel_operations = set() + self._cancel_pending_operations.update(cancel_operations) + self._cancel_pending_endpoints.update( + endpoint for endpoint, _epoch in cancel_operations + ) + orphaned = [ + (self.disagg_request_id, pending_task.slice_id) + for pending_task in self._kv_tasks + if pending_task.status == TaskStatus.TRANSFERRING or pending_task is task + ] + return cancel_endpoints, cancel_operations, orphaned def receive(self, slice: KVSlice) -> None: - if self.transfer_start_time is None: - self.transfer_start_time = tensorrt_llm.bindings.global_steady_clock_now() - params = self._base_args.params - slice_id = len(self._kv_tasks) - task = KVRecvTask( - self.disagg_request_id, - slice, - slice_id, - params, - aux_slot=self.aux_slot, - ) - self._kv_tasks.append(task) + with self.lock: + params = self._base_args.params + slice_id = len(self._kv_tasks) + task = KVRecvTask( + self.disagg_request_id, + slice, + slice_id, + params, + aux_slot=self.aux_slot, + ) + self._kv_tasks.append(task) + if self._sealed or self._closed: + task.fail( + RuntimeError(f"RxSession {self.disagg_request_id} is sealed; receive rejected") + ) + return + if self.transfer_start_time is None: + self.transfer_start_time = tensorrt_llm.bindings.global_steady_clock_now() self._receiver.dispatch_task(task) def process_kv_agent_result( @@ -1889,136 +3795,286 @@ def process_kv_agent_result( sizes=None, src_base=None, transfer_size: int = 0, + sender_endpoint: Optional[str] = None, + request_epoch: Optional[int] = None, ): + operation_key = ("kv", sender_slice_id, peer_rank) + bounce_action: Optional[Callable[[], None]] = None + protocol_error: Optional[RuntimeError] = None + cancel_operations: set[tuple[str, int]] = set() + task: Optional[KVRecvTask] = None with self.lock: - self.kv_cache_size_bytes += transfer_size - assert sender_slice_id < len(self._kv_tasks), ( - f"Receiver got slice_id={sender_slice_id} from sender but only has " - f"{len(self._kv_tasks)} receive task(s) for request {self.request_id}. " - f"Sender/receiver slice count mismatch." - ) - task = self._kv_tasks[sender_slice_id] - if status == AgentResult.SUCCESS: - from .bounce import scatter_write_result - - on_done = None - if is_last_slice: + try: + validation = self._validate_receive_operation_unlocked( + operation_key, + sender_endpoint=sender_endpoint, + request_epoch=request_epoch, + ) + if validation == "duplicate": + return + if not 0 <= sender_slice_id < len(self._kv_tasks): + raise RuntimeError( + f"result names invalid slice {sender_slice_id}; receiver has " + f"{len(self._kv_tasks)} slice(s) for request {self.request_id}" + ) + except RuntimeError as error: + protocol_error = error + cancel_operations = self._fail_protocol_unlocked(error) + if protocol_error is None: + self.kv_cache_size_bytes += transfer_size + task = self._kv_tasks[sender_slice_id] + rid_slice = (self.disagg_request_id, task.slice_id) + complete_task = False + if status == AgentResult.SUCCESS and is_last_slice: task.last_slice_count += 1 - if task.last_slice_count == task.expected_transfers: - # Completing message: defer task.complete()+perf until the scatter has actually - # landed. scatter_write_result fires this inline for the non-bounced path, or on - # the scatter worker (after cudaStreamSynchronize) for the bounced path, so the - # gen consumer never observes completion before the KV is scattered into place. - request_id = self.request_id - ri = self._receiver._registrar.self_rank_info - instance_name, instance_rank = ri.instance_name, ri.instance_rank - - def on_done( - success, - task=task, + complete_task = task.last_slice_count == task.expected_transfers + elif status == AgentResult.FAILED: + detail = ( + f"KV transfer failed for request {self.request_id} " + f"slice={sender_slice_id} peer_rank={peer_rank} " + f"is_last_slice={is_last_slice} (reported by remote agent; " + "see sender-side log for nixl_status)" + ) + logger.error(detail) + if self._terminal_status != SessionStatus.CANCELLED: + task.fail(RuntimeError(detail)) + self._terminal_status = SessionStatus.ERROR + else: + if status != AgentResult.SUCCESS: + protocol_error = RuntimeError( + f"unknown KV result status {status!r} for request {self.request_id}" + ) + cancel_operations = self._fail_protocol_unlocked(protocol_error) + + if protocol_error is None: + on_done = None + if complete_task: + on_done = self._make_kv_settlement_callback( + task, peer_rank=peer_rank, - sender_slice_id=sender_slice_id, - request_id=request_id, - 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. - if not success: - task.fail( - RuntimeError( - f"KV bounce scatter failed for request {request_id} " - f"slice={sender_slice_id}" - ) - ) - return - if task.status == TaskStatus.ERROR: - return # a concurrent FAILED writer already failed it; don't un-fail - try: - if task._perf_timer is not None: - task._perf_timer.record_task_end(peer_rank) - task.print_perf_info(peer_rank, instance_name, instance_rank) - except Exception as e: # perf is best-effort; never block completion - logger.warning( - f"KV transfer perf logging failed for request {request_id} " - f"slice={sender_slice_id}: {e}" - ) - 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. - 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( - f"KV transfer complete for request {request_id} " - f"slice={sender_slice_id}" + complete_task=complete_task, + settle_bounce=False, + ) + if status == AgentResult.SUCCESS: + from .bounce import scatter_write_result + + def record_success() -> None: + scatter_write_result( + self._receiver._bounce, + rid_slice, + peer_rank, + dst_ptrs, + sizes, + src_base, + on_done, ) - scatter_write_result( - self._receiver._bounce, - (self.disagg_request_id, task.slice_id), - peer_rank, - dst_ptrs, - sizes, - src_base, - on_done, - ) - elif status == AgentResult.FAILED: - detail = ( - f"KV transfer failed for request {self.request_id} slice={sender_slice_id} " - f"peer_rank={peer_rank} is_last_slice={is_last_slice} " - f"(reported by remote agent; see sender-side log for nixl_status)" - ) - logger.error(detail) - # Drain-before-release: record this writer FAILED; the owner frees the shared region - # only once every fan-in writer is terminal (freeing now could race a sibling's RMA). - self._receiver._bounce.record_failure( - (self.disagg_request_id, task.slice_id), peer_rank - ) - task.fail(RuntimeError(detail)) - if self._terminal_status is None: # Don't overwrite CANCELLED with ERROR - self._terminal_status = SessionStatus.ERROR - else: - raise ValueError( - f"Session {self.request_id} received unknown task status: {status.value}" + bounce_action = record_success + else: + + def record_failure() -> None: + self._receiver._bounce.record_failure( + rid_slice, peer_rank, on_done=on_done + ) + + bounce_action = record_failure + + if protocol_error is not None: + self._publish_protocol_failure(protocol_error, cancel_operations) + return + + try: + if bounce_action is not None: + bounce_action() + except Exception as error: + with self.lock: + protocol_error = RuntimeError( + f"bounce result validation/settlement failed for request " + f"{self.request_id} slice={sender_slice_id} rank={peer_rank}: {error}" ) + if task is not None: + task.fail(protocol_error) + cancel_operations = self._fail_protocol_unlocked(protocol_error) + self._publish_protocol_failure(protocol_error, cancel_operations) + return + + with self.lock: + self._retire_receive_operation_unlocked(operation_key) + finalize_close = self._should_finalize_close_unlocked() + if finalize_close: + self._finalize_close() - def process_aux_agent_result(self, _peer_rank: int, status: AgentResult): + def process_aux_agent_result( + self, + peer_rank: int, + status: AgentResult, + *, + sender_endpoint: Optional[str] = None, + request_epoch: Optional[int] = None, + ): # Aux is session-level (not per-slice); expected_transfers is identical # across all kv_tasks, so any task provides the right count. + operation_key = ("aux", peer_rank) + finalize_close = False + protocol_error: Optional[RuntimeError] = None + cancel_operations: set[tuple[str, int]] = set() with self.lock: - if not self._kv_tasks: - logger.warning( - f"Aux result received before any KV tasks for request {self.request_id}" + try: + validation = self._validate_receive_operation_unlocked( + operation_key, + sender_endpoint=sender_endpoint, + request_epoch=request_epoch, ) + except RuntimeError as error: + protocol_error = error + cancel_operations = self._fail_protocol_unlocked(error) + validation = "reject" + if validation == "duplicate": return - task = self._kv_tasks[0] - if status == AgentResult.SUCCESS: - self._aux_count += 1 - - if self._aux_count == task.expected_transfers: - self._aux_status = TaskStatus.TRANSFERRED - elif self._aux_count > task.expected_transfers: + if protocol_error is None and not self._kv_tasks: + protocol_error = RuntimeError( + f"aux result arrived before KV dispatch for request {self.request_id}" + ) + cancel_operations = self._fail_protocol_unlocked(protocol_error) + if protocol_error is None and self._terminal_status != SessionStatus.CANCELLED: + task = self._kv_tasks[0] + if status == AgentResult.SUCCESS: + self._aux_count += 1 + + if self._aux_count == task.expected_transfers: + self._aux_status = TaskStatus.TRANSFERRED + elif self._aux_count > task.expected_transfers: + self._aux_status = TaskStatus.ERROR + self._exception = RuntimeError( + f"Session {self.request_id} received too many aux transfers" + ) + 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} received too many aux transfers" - ) + self._exception = RuntimeError(f"Session {self.request_id} aux transfer failed") 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") - if self._terminal_status is None: - self._terminal_status = SessionStatus.ERROR - else: - raise ValueError( - f"Session {self.request_id} received unknown aux send status: {status}" + else: + protocol_error = RuntimeError( + f"unknown aux result status {status!r} for request {self.request_id}" + ) + cancel_operations = self._fail_protocol_unlocked(protocol_error) + if protocol_error is None: + self._retire_receive_operation_unlocked(operation_key) + finalize_close = self._should_finalize_close_unlocked() + if protocol_error is not None: + self._publish_protocol_failure(protocol_error, cancel_operations) + return + if finalize_close: + self._finalize_close() + + def _make_kv_settlement_callback( + self, + task: KVRecvTask, + *, + peer_rank: int, + complete_task: bool, + settle_bounce: bool, + ) -> Callable[[bool], None]: + request_id = self.request_id + sender_slice_id = task.slice_id + ri = self._receiver._registrar.self_rank_info + instance_name, instance_rank = ri.instance_name, ri.instance_rank + + def settled(success: bool) -> None: + try: + if not success: + if task.status != TaskStatus.ERROR: + task.fail( + RuntimeError( + f"KV bounce settlement failed for request {request_id} " + f"slice={sender_slice_id}" + ) + ) + return + if not complete_task or task.status == TaskStatus.ERROR: + return + try: + if task._perf_timer is not None: + task._perf_timer.record_task_end(peer_rank) + task.print_perf_info(peer_rank, instance_name, instance_rank) + except Exception as error: + logger.warning( + f"KV transfer perf logging failed for request {request_id} " + f"slice={sender_slice_id}: {error}" + ) + task.complete() + 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( + f"KV transfer complete for request {request_id} slice={sender_slice_id}" ) + finally: + if settle_bounce: + self._retire_bounce_settlement(sender_slice_id) + + return settled + + def _make_bounce_settlement_callback(self, task: KVRecvTask) -> Callable[[bool], None]: + """Retire one slice's arena credit on every proven settlement path.""" + request_id = self.request_id + slice_id = task.slice_id + + def settled(success: bool) -> None: + try: + if not success and task.status != TaskStatus.ERROR: + task.fail( + RuntimeError( + f"KV bounce settlement failed for request {request_id} slice={slice_id}" + ) + ) + finally: + self._retire_bounce_settlement(slice_id) + + return settled + + def _retire_bounce_settlement(self, slice_id: int) -> None: + finalize_close = False + with self.lock: + if slice_id not in self._pending_bounce_settlements: + return + self._pending_bounce_settlements.remove(slice_id) + self._pending_scatter_callbacks = len(self._pending_bounce_settlements) + finalize_close = self._should_finalize_close_unlocked() + if finalize_close: + self._finalize_close() + + def _publish_protocol_failure( + self, + error: RuntimeError, + cancel_operations: set[tuple[str, int]], + ) -> None: + self._receiver._fail_protocol(error) + with self.lock: + bounced_slices = list(self._pending_bounce_settlements) + for slice_id in bounced_slices: + try: + self._receiver._bounce.orphan_reservation((self.disagg_request_id, slice_id)) + except Exception as bounce_error: + logger.error( + "Failed to retain bounce reservation after protocol error for " + "request %s slice=%s: %s", + self.request_id, + slice_id, + bounce_error, + ) + if cancel_operations: + endpoints = {endpoint for endpoint, _epoch in cancel_operations} + try: + self._receiver.send_cancel_to_senders( + self.disagg_request_id, endpoints, cancel_operations + ) + except Exception: + # send_cancel_to_senders already records a receiver-wide + # fail-stop state. Lifetime credits intentionally remain live. + pass @property def exception(self) -> Optional[Exception]: @@ -2054,6 +4110,127 @@ def has_failed(self) -> bool: """Non-blocking check: has the transfer failed or been cancelled?""" return self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED) + def _retire_receive_operation_unlocked(self, operation_key: tuple) -> None: + if operation_key in self._retired_operation_keys: + return + if self._outstanding_operations <= 0: + raise RuntimeError(f"RxSession {self.disagg_request_id} operation credit underflow") + self._retired_operation_keys.add(operation_key) + self._outstanding_operations -= 1 + + def _retire_scatter_callback(self) -> None: + finalize_close = False + with self.lock: + if self._pending_scatter_callbacks <= 0: + raise RuntimeError(f"RxSession {self.disagg_request_id} scatter credit underflow") + self._pending_scatter_callbacks -= 1 + finalize_close = self._should_finalize_close_unlocked() + if finalize_close: + self._finalize_close() + + def process_cancel_ack(self, sender_endpoint: str, request_epoch: Optional[int] = None) -> None: + finalize_close = False + drained_bounces: list[tuple[int, int]] = [] + protocol_error: Optional[RuntimeError] = None + with self.lock: + if request_epoch is None: + if self._v2_enabled: + protocol_error = RuntimeError( + f"legacy cancel ACK received for v2 request {self.disagg_request_id}" + ) + elif sender_endpoint not in self._cancel_pending_endpoints: + return + else: + self._cancel_pending_endpoints.discard(sender_endpoint) + if ( + self._terminal_status in (SessionStatus.CANCELLED, SessionStatus.ERROR) + and not self._cancel_pending_endpoints + ): + self._outstanding_operations = 0 + drained_bounces = [ + (self.disagg_request_id, task.slice_id) for task in self._kv_tasks + ] + else: + operation = (sender_endpoint, request_epoch) + if operation not in self._cancel_pending_operations: + known = any( + expected.sender_endpoint == sender_endpoint + and expected.request_epoch == request_epoch + for expected in self._expected_operations.values() + ) + if known: + return + protocol_error = RuntimeError( + "cancel ACK does not match an active request operation: " + f"rid={self.disagg_request_id} endpoint={sender_endpoint} " + f"epoch={request_epoch}" + ) + else: + self._cancel_pending_operations.discard(operation) + for operation_key, expected in self._expected_operations.items(): + if ( + operation_key not in self._retired_operation_keys + and expected.sender_endpoint == sender_endpoint + and expected.request_epoch == request_epoch + ): + self._retire_receive_operation_unlocked(operation_key) + if not any( + endpoint == sender_endpoint + for endpoint, _epoch in self._cancel_pending_operations + ): + self._cancel_pending_endpoints.discard(sender_endpoint) + for slice_id in list(self._pending_bounce_settlements): + has_live_writer = any( + key[0] == "kv" + and key[1] == slice_id + and key not in self._retired_operation_keys + for key in self._expected_operations + ) + if not has_live_writer: + drained_bounces.append((self.disagg_request_id, slice_id)) + if protocol_error is not None: + self._fail_protocol_unlocked(protocol_error) + for rid_slice in drained_bounces: + self._receiver._bounce.confirm_drained(rid_slice) + if protocol_error is not None: + self._publish_protocol_failure(protocol_error, set()) + return + with self.lock: + finalize_close = self._should_finalize_close_unlocked() + if finalize_close: + self._finalize_close() + + def is_cancel_pending(self, sender_endpoint: str, request_epoch: int) -> bool: + with self.lock: + return (sender_endpoint, request_epoch) in self._cancel_pending_operations + + def validate_remote_cancel(self, sender_endpoint: Optional[str], request_epoch: int) -> None: + with self.lock: + if any( + expected.sender_endpoint == sender_endpoint + and expected.request_epoch == request_epoch + for expected in self._expected_operations.values() + ): + return + error = RuntimeError( + "remote cancel does not match an active request incarnation: " + f"rid={self.disagg_request_id} endpoint={sender_endpoint} " + f"epoch={request_epoch}" + ) + self._fail_protocol_unlocked(error) + raise error + + def _is_drained_unlocked(self) -> bool: + return ( + self._outstanding_operations == 0 + and self._pending_scatter_callbacks == 0 + and not self._cancel_pending_endpoints + and not self._cancel_pending_operations + ) + + def _should_finalize_close_unlocked(self) -> bool: + return self._close_requested and not self._closed and self._is_drained_unlocked() + def cancel(self) -> None: """Cancel the session and notify the remote sender. @@ -2061,10 +4238,24 @@ def cancel(self) -> None: Only INIT tasks have their events signalled immediately. The lock serializes with process_kv_agent_result() / process_aux_agent_result(). """ + sender_endpoints: set[str] + orphaned: list[tuple[int, int]] = [] with self.lock: if self._terminal_status == SessionStatus.CANCELLED: return + self._sealed = True self._terminal_status = SessionStatus.CANCELLED + sender_endpoints = set(self._sender_endpoints) + cancel_operations = { + (expected.sender_endpoint, expected.request_epoch) + for operation_key, expected in self._expected_operations.items() + if operation_key not in self._retired_operation_keys + and expected.request_epoch is not None + } + self._cancel_pending_operations.update(cancel_operations) + self._cancel_pending_endpoints.update( + endpoint for endpoint, _epoch in cancel_operations + ) exc = RuntimeError(f"RxSession {self.disagg_request_id} cancelled") for task in self._kv_tasks: rid_slice = (self.disagg_request_id, task.slice_id) @@ -2074,18 +4265,32 @@ def cancel(self) -> None: self._receiver._bounce.release_idle_reservation(rid_slice) task.fail(exc) elif task.status == TaskStatus.TRANSFERRING: - # A write may still be mid-flight, so quarantine the region rather than freeing - # it; this keeps a cancelled transfer from leaking. No-op when bounce is off. - self._receiver._bounce.orphan_reservation(rid_slice) + orphaned.append(rid_slice) + for rid_slice in orphaned: + self._receiver._bounce.orphan_reservation(rid_slice) # Send outside the lock to avoid holding it during I/O. - self._receiver.send_cancel_to_senders(self.disagg_request_id, self._sender_endpoints) + self._receiver.send_cancel_to_senders( + self.disagg_request_id, + sender_endpoints, + cancel_operations, + ) def has_transferring_tasks(self) -> bool: """True if any KV task is currently mid-write (TRANSFERRING). cancel_request() must return False while this is True. """ - return any(t.status == TaskStatus.TRANSFERRING for t in self._kv_tasks) + with self.lock: + return self._has_transferring_tasks_unlocked() + + def _has_transferring_tasks_unlocked(self) -> bool: + return not self._is_drained_unlocked() + + def seal_and_check_quiescent(self) -> bool: + """Atomically prevent new receive work and check whether writes drained.""" + with self.lock: + self._sealed = True + return not self._has_transferring_tasks_unlocked() def wait_complete(self, blocking: bool = False) -> Optional[WaitResult]: """Poll or block until transfer completes. @@ -2126,9 +4331,26 @@ def wait_complete(self, blocking: bool = False) -> Optional[WaitResult]: return WaitResult.COMPLETED def close(self): - if getattr(self, "_closed", False): + with self.lock: + if self._closed: + return + # Publish the strong drain reference while the closed-state check + # is still serialized with finalization. Otherwise two concurrent + # close() calls can let the second reinsert an already-cleared + # session into Receiver._draining_sessions forever. + self._receiver.retain_draining_session(self) + self._sealed = True + self._close_requested = True + finalize_close = self._should_finalize_close_unlocked() + if not finalize_close: return - self._closed = True + self._finalize_close() + + def _finalize_close(self) -> None: + with self.lock: + if self._closed or not self._is_drained_unlocked(): + return + self._closed = True if self._aux_buffer is not None and self.aux_slot is not None: self._aux_buffer.free_slot(self.aux_slot) self.aux_slot = None @@ -2138,7 +4360,10 @@ def close(self): # leaked; a no-op for finished or non-bounce transfers. for task in self._kv_tasks: self._receiver._bounce.orphan_reservation((self.disagg_request_id, task.slice_id)) - self._receiver.clear_session(self.disagg_request_id) + self._receiver.clear_session( + self.disagg_request_id, + self.request_epoch if self._v2_enabled else None, + ) def __enter__(self): return self @@ -2194,7 +4419,10 @@ def handle_message(messages: list[bytes]) -> bool: self._messenger.start_listener(handle_message) def _handle_rank_info_request(self, send_id: bytes, _message: list[bytes]): - self._messenger.send([send_id, self._rank_info.to_bytes()]) + # The extra frame is backward compatible: legacy receivers consume + # the RankInfo frame and ignore the remainder. New receivers treat a + # missing frame as protocol v1 and keep the legacy retention path. + self._messenger.send([send_id, self._rank_info.to_bytes(), _encode_protocol_capabilities()]) def __del__(self): try: @@ -2246,6 +4474,11 @@ class TransferWorkerConfig: class TransferWorker: def __init__(self, config: TransferWorkerConfig): + self._shutdown = False + self._shutdown_lock = threading.Lock() + self._shutdown_complete = threading.Event() + self._shutdown_thread: Optional[threading.Thread] = None + self._shutdown_error: Optional[BaseException] = None self._config = config kvm = config.kv_cache_manager self._aux_buffer = _make_aux_buffer( @@ -2294,6 +4527,12 @@ def create_rx_session(self, request: LlmRequest) -> RxSession: def has_all_peer_req_infos_for_send(self, unique_rid: int) -> bool: return self._sender.has_all_peer_req_infos(unique_rid) + def pin_peer_req_infos_for_send(self, unique_rid: int) -> None: + self._sender.pin_peer_req_infos(unique_rid) + + def cancel_peer_req_infos_for_send(self, unique_rid: int) -> bool: + return self._sender.cancel_peer_req_infos(unique_rid) + def sweep_stale_req_infos(self): """Forward to Sender to evict orphaned RecvReqInfo from ADP broadcast.""" self._sender.sweep_stale_req_infos() @@ -2373,10 +4612,80 @@ def page_table(self): assert self._rank_info is not None return self._rank_info.page_table - def shutdown(self): - if getattr(self, "_shutdown", False): - return - self._shutdown = True + def shutdown(self) -> Optional[threading.Event]: + if self._is_progress_thread(): + return self._defer_shutdown() + run_shutdown = False + with self._shutdown_lock: + if self._shutdown_complete.is_set(): + self._raise_shutdown_error() + return + if not self._shutdown: + self._shutdown = True + run_shutdown = True + + if run_shutdown: + self._run_shutdown() + else: + self._shutdown_complete.wait() + self._raise_shutdown_error() + return None + + def _is_progress_thread(self) -> bool: + """Whether teardown would join or wait on the current internal thread.""" + current = threading.current_thread() + receiver = getattr(self, "_receiver", None) + receiver_messenger = getattr(receiver, "_messenger", None) + if current is getattr(receiver_messenger, "_listener_thread", None): + return True + bounce = getattr(self, "_bounce", None) + if current is getattr(bounce, "_scatter_thread", None): + return True + sender = getattr(self, "_sender", None) + sender_messenger = getattr(sender, "_messenger", None) + if current is getattr(sender_messenger, "_listener_thread", None): + return True + if current in getattr(sender, "_worker_threads", ()): + return True + rank_info_server = getattr(self, "_rank_info_server", None) + rank_info_messenger = getattr(rank_info_server, "_messenger", None) + return current is getattr(rank_info_messenger, "_listener_thread", None) + + def _defer_shutdown(self) -> threading.Event: + """Allow the current progress callback to return before teardown waits on it.""" + with self._shutdown_lock: + if self._shutdown_complete.is_set() or self._shutdown: + return self._shutdown_complete + self._shutdown = True + self._shutdown_thread = threading.Thread( + target=self._run_shutdown, + name="trtllm-transfer-worker-shutdown", + daemon=True, + ) + shutdown_thread = self._shutdown_thread + logger.warning( + "TransferWorker shutdown was requested from an internal progress thread; " + "deferring teardown until the callback returns" + ) + shutdown_thread.start() + return self._shutdown_complete + + def _run_shutdown(self) -> None: + try: + self._shutdown_resources() + except BaseException as error: + self._shutdown_error = error + logger.exception("TransferWorker shutdown stopped before memory teardown was safe") + finally: + self._shutdown_complete.set() + + def _raise_shutdown_error(self) -> None: + if self._shutdown_error is not None: + raise RuntimeError( + "TransferWorker shutdown did not complete safely" + ) from self._shutdown_error + + def _shutdown_resources(self) -> None: # Use getattr guards: __init__ may have failed partway, leaving some # attributes unset. Without them, __del__ -> shutdown() raises # AttributeError and ZMQ resources from already-created sub-objects @@ -2396,10 +4705,10 @@ def shutdown(self): # buffers leak until process exit. bounce = getattr(self, "_bounce", None) if bounce is not None: - try: - bounce.close() - except Exception as e: - logger.warning(f"TransferWorker.shutdown: bounce close failed: {e}") + # A close failure means the scatter thread may still touch its VMM + # region. Propagate immediately and retain every later registered + # memory object instead of invalidating memory under a live thread. + bounce.close() # Deregister NIXL memory before agent.shutdown so pinned GPU memory is released # (e.g. when the KV cache manager is recreated after profiling). agent = getattr(self, "_agent", None) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 2c257a39f731..abfb45aac3cf 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -1,7 +1,24 @@ +# 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. + +import hashlib import os +import threading import time import uuid -from collections import defaultdict +from collections import OrderedDict, defaultdict from itertools import chain from typing import Any, Callable, Dict, List, Optional, cast @@ -10,6 +27,14 @@ import tensorrt_llm.bindings from tensorrt_llm import logger +from tensorrt_llm._torch.disaggregation.async_consensus import ( + PROTOCOL_VERSION, + AsyncConsensusCoordinator, + ConsensusEvent, + ConsensusEventKind, + ConsensusOutcome, + MpiConsensusTransport, +) from tensorrt_llm._torch.disaggregation.base.transfer import ( KVSlice, RxSessionBase, @@ -30,7 +55,7 @@ ) from tensorrt_llm._torch.disaggregation.resource.page import MambaLayerGroup from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool -from tensorrt_llm._torch.distributed.communicator import Distributed +from tensorrt_llm._torch.distributed.communicator import Distributed, MPIDist from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import KvCacheTransceiver from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager @@ -42,6 +67,16 @@ from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig from tensorrt_llm.mapping import Mapping +_ASYNC_TERMINAL_ENV = "TRTLLM_PYTHON_TRANSCEIVER_ASYNC_CTX_TERMINAL_CONSENSUS" +_ASYNC_PEER_READY_ENV = "TRTLLM_PYTHON_TRANSCEIVER_ASYNC_CTX_PEER_READY_CONSENSUS" +_CONTEXT_ACTIVATION_DIGEST_ENV = "TRTLLM_PYTHON_TRANSCEIVER_CONTEXT_ACTIVATION_DIGEST" +_ASYNC_STARTUP_TAG = "TRTLLM_PYTHON_TRANSCEIVER_ASYNC_CONSENSUS" +_ASYNC_READY_CANCELLED_EPOCH_ATTR = "_trtllm_async_ready_cancelled_epoch" +_MAX_RETIRED_CONSENSUS_REQUESTS = 65536 +_STARTUP_ROLLBACK_TIMEOUT_S = 30.0 +_CONSENSUS_STARTUP_CLOSE_TIMEOUT_S = 1.0 +_ASYNC_READY_MAX_IDLE_SLEEP_S = 0.01 + def _find_consensus_request_ids(request_ids_all_ranks, sync_size): frequency_map = defaultdict(int) @@ -63,6 +98,7 @@ def __init__( dist: Distributed, kv_cache_manager: KVCacheManager, cache_transceiver_config: CacheTransceiverConfig, + publish_disaggregated_params: bool = True, ): self._dist: Distributed = dist self._kv_cache_manager = kv_cache_manager @@ -72,9 +108,35 @@ def __init__( self._sender_future_timeout_ms = ( cache_transceiver_config.kv_transfer_sender_future_timeout_ms ) + self._publish_disaggregated_params = publish_disaggregated_params + # Initialize optional consensus state without adding a collective to + # the default-off startup path. Explicit opt-in is negotiated later by + # piggybacking on the existing endpoint exchange. + self._init_async_consensus(cache_transceiver_config) self._check_compatible() + self._init_sync_policy() self._reuse_adapter: CacheReuseAdapter = create_cache_reuse_adapter(kv_cache_manager) + # Initialize teardown-visible ownership before native workers start. + # Explicit consensus startup happens after TransferWorker construction, + # so constructor rollback must be able to use the normal shutdown path + # even when communicator negotiation fails partway through __init__. + self._send_sessions: Dict[int, TxSessionBase] = {} + self._recv_sessions: Dict[int, RxSessionBase] = {} + self._send_reqs = {} + self._recv_reqs = {} + self._wait_reqs = {} + self._legacy_failed_sessions: set[int] = set() + self._shutdown = False + self._shutdown_complete = False + self._shutdown_metadata_leases_complete = False + self._shutdown_sessions_complete = False + self._shutdown_consensus_complete = False + self._shutdown_worker_complete = False + self._shutdown_worker_event: Optional[threading.Event] = None + self._shutdown_deferred_errors: list[tuple[str, Exception]] = [] + self._async_ready_idle_wakeup = threading.Event() + self._device_id = torch.cuda.current_device() logger.info(f"device_id: {self._device_id} in KvCacheTransceiverV2") self._instance_name = self._broadcast_instance_name() @@ -94,15 +156,13 @@ def __init__( ) ) self._dp_rank = mapping.tp_rank if mapping.enable_attention_dp else 0 - self._context_info_endpoint = self._broadcast_context_endpoint() - self._init_sync_policy() - self._exchange_rank_info() + try: + self._context_info_endpoint = self._broadcast_context_endpoint() + self._exchange_rank_info() + except Exception: + self._rollback_failed_startup() + raise - self._send_sessions: Dict[int, TxSessionBase] = {} - self._recv_sessions: Dict[int, RxSessionBase] = {} - self._send_reqs = {} - self._recv_reqs = {} - self._wait_reqs = {} self._page_table = self._transfer_worker.page_table # _slice_num_bytes() is this rank's KV shard, so scale by tp_size to get the request total (kv_cache_size), # except under attention DP where the local count already is the total. @@ -113,6 +173,240 @@ def __init__( self._ever_had_send_session: bool = False self._ever_had_recv_session: bool = False + @staticmethod + def _parse_binary_env(name: str, value: str) -> bool: + if value not in ("0", "1"): + raise ValueError(f"{name} must be 0 or 1, got {value!r}") + return value == "1" + + def _rollback_failed_startup(self) -> None: + """Release native and consensus resources after constructor failure.""" + try: + worker_event = self.shutdown() + if isinstance(worker_event, threading.Event): + if not worker_event.wait(_STARTUP_ROLLBACK_TIMEOUT_S): + logger.error( + "Python transceiver startup rollback timed out waiting " + "for deferred native shutdown" + ) + return + self.shutdown() + except Exception as error: + # Preserve the startup exception. shutdown() already attempts every + # independent teardown step before surfacing its first error. + logger.error(f"Python transceiver startup rollback failed: {error}") + + def _init_context_activation_digest(self) -> None: + enabled = self._parse_binary_env( + _CONTEXT_ACTIVATION_DIGEST_ENV, + os.getenv(_CONTEXT_ACTIVATION_DIGEST_ENV, "0"), + ) + self._context_activation_digest = hashlib.sha256() if enabled else None + self._context_activation_count = 0 + self._context_activation_digest_logged = False + + def _record_context_activation_ids(self, request_ids: List[int]) -> None: + """Record an order-sensitive digest without logging request IDs.""" + digest = getattr(self, "_context_activation_digest", None) + if digest is None: + return + for request_id in request_ids: + payload = str(int(request_id)).encode("ascii") + digest.update(len(payload).to_bytes(4, byteorder="big")) + digest.update(payload) + self._context_activation_count += len(request_ids) + + def _log_context_activation_digest(self) -> None: + digest = getattr(self, "_context_activation_digest", None) + if digest is None or getattr(self, "_context_activation_digest_logged", False): + return + logger.info( + "PYTHON_CONTEXT_ACTIVATION_SEQUENCE " + f"rank={self._dist.rank} " + f"count={self._context_activation_count} " + f"digest={digest.hexdigest()} " + "algorithm=sha256-length-prefixed-decimal-v1" + ) + self._context_activation_digest_logged = True + + def _init_async_consensus(self, cache_transceiver_config: CacheTransceiverConfig) -> None: + self._init_context_activation_digest() + terminal_value = os.getenv(_ASYNC_TERMINAL_ENV, "0") + peer_ready_value = os.getenv(_ASYNC_PEER_READY_ENV, "0") + self._async_terminal_flag_value = terminal_value + self._async_peer_ready_flag_value = peer_ready_value + self._async_consensus_config = cache_transceiver_config + + self._async_terminal_consensus_enabled = False + self._async_peer_ready_consensus_enabled = False + self._async_consensus: Optional[AsyncConsensusCoordinator] = None + self._async_terminal_epoch: OrderedDict[int, int] = OrderedDict() + self._async_terminal_published: Dict[int, int] = {} + self._async_terminal_commits: Dict[int, ConsensusEvent] = {} + self._async_terminal_cancelled: Dict[int, tuple[int, LlmRequest]] = {} + self._context_cancelled_request_ids: list[int] = [] + self._async_ready_epoch: OrderedDict[int, int] = OrderedDict() + self._async_ready_published: Dict[int, int] = {} + self._async_ready_prepared: Dict[tuple[int, int], LlmRequest] = {} + self._async_ready_released: set[tuple[int, int]] = set() + self._async_ready_activated: Dict[tuple[int, int], LlmRequest] = {} + self._async_ready_metadata_leases: set[int] = set() + self._async_ready_acknowledged: set[tuple[int, int]] = set() + self._async_ready_withdrawn: set[tuple[int, int]] = set() + # READY_ABORT is intentionally broadcast to every participant, even + # when a rank has not materialized the request locally yet. Keep a + # request-less tombstone until the request arrives (or until a bounded + # post-finalize replay window consumes it); acknowledging the abort + # must not depend on local request timing. + self._async_ready_aborted: Dict[tuple[int, int], Optional[LlmRequest]] = {} + self._async_ready_finalized_without_request: OrderedDict[int, int] = OrderedDict() + self._async_consensus_counters: Dict[str, int] = defaultdict(int) + + # A singleton has no cross-rank state to reconcile. Parse locally and + # treat either opt-in as a no-op instead of constructing a communicator. + if self._mapping.world_size == 1: + terminal_requested = self._parse_binary_env(_ASYNC_TERMINAL_ENV, terminal_value) + peer_ready_requested = self._parse_binary_env(_ASYNC_PEER_READY_ENV, peer_ready_value) + if terminal_requested or peer_ready_requested: + logger.info( + "PYTHON_ASYNC_CONSENSUS transition=singleton_noop " + f"terminal={int(terminal_requested)} " + f"peer_ready={int(peer_ready_requested)}" + ) + return + + # Multi-rank parsing and qualification happen in _exchange_rank_info. + # A malformed or mismatched same-version opt-in must reach that + # universally-entered allgather before every new worker rejects it. + + def _async_startup_descriptor(self) -> tuple: + config = self._async_consensus_config + topology = ( + int(self._mapping.world_size), + int(self._mapping.tp_size), + int(self._mapping.pp_size), + int(self._mapping.cp_size), + bool(self._mapping.enable_attention_dp), + ) + distributed_runtime = ( + "MPI" + if isinstance(self._dist, MPIDist) + else f"{type(self._dist).__module__}.{type(self._dist).__qualname__}" + ) + return ( + PROTOCOL_VERSION, + str(config.backend), + str(config.transceiver_runtime), + distributed_runtime, + topology, + self._async_terminal_flag_value, + self._async_peer_ready_flag_value, + ) + + def _complete_async_consensus_startup(self, gathered: list) -> list[str]: + """Validate opt-in metadata piggybacked on endpoint exchange. + + With both flags off, every contribution remains the exact legacy + endpoint string. Explicit opt-in intentionally requires a same-version + worker group; it is not a rolling-upgrade compatibility mechanism. + """ + + def is_tagged(value: Any) -> bool: + return isinstance(value, tuple) and len(value) == 3 and value[0] == _ASYNC_STARTUP_TAG + + tagged = [is_tagged(value) for value in gathered] + if not any(tagged): + return [cast(str, endpoint) for endpoint in gathered] + if not all(tagged): + raise RuntimeError( + "asynchronous Python consensus explicit opt-in requires a " + "same-version worker group; mixed tagged and legacy endpoint " + f"contributions were gathered: {gathered}" + ) + + endpoints = [cast(str, value[1]) for value in gathered] + descriptors = [tuple(value[2]) for value in gathered] + descriptor = self._async_startup_descriptor() + if len(descriptors) != self._mapping.world_size or any( + peer_descriptor != descriptor for peer_descriptor in descriptors + ): + raise RuntimeError( + "asynchronous Python consensus startup descriptor mismatch " + f"across worker ranks: local={descriptor}, gathered={descriptors}" + ) + + terminal_requested = self._parse_binary_env( + _ASYNC_TERMINAL_ENV, self._async_terminal_flag_value + ) + peer_ready_requested = self._parse_binary_env( + _ASYNC_PEER_READY_ENV, self._async_peer_ready_flag_value + ) + requested = terminal_requested or peer_ready_requested + if not requested: + raise RuntimeError( + "asynchronous Python consensus startup metadata was tagged without an enabled mode" + ) + + negotiation_domain = ( + isinstance(self._dist, MPIDist) + and self._mapping.tp_size == 1 + and self._mapping.cp_size == 1 + and not self._mapping.enable_attention_dp + and self._mapping.pp_size > 1 + and self._mapping.world_size == self._mapping.pp_size + ) + if not negotiation_domain: + raise RuntimeError( + "asynchronous Python transceiver CTX consensus currently requires " + "the MPI distributed runtime, TP1, CP1, non-ADP, PP>1, and a " + "PP domain equal to the worker world" + ) + + config = self._async_consensus_config + qualified = config.backend == "NIXL" and config.transceiver_runtime == "PYTHON" + if not qualified: + raise RuntimeError( + "asynchronous Python transceiver CTX consensus currently requires " + "backend='NIXL' and transceiver_runtime='PYTHON'" + ) + + self._async_terminal_consensus_enabled = terminal_requested + self._async_peer_ready_consensus_enabled = peer_ready_requested + + participants = tuple(int(rank) for rank in self._mapping.pp_group) + # Construction runs on the executor thread after the existing startup + # collectives. TransferWorker's background threads progress NIXL, not + # this Distributed MPI communicator; the transport duplicates the + # communicator before its asynchronous point-to-point traffic starts. + transport = MpiConsensusTransport(participants) + try: + coordinator = AsyncConsensusCoordinator( + transport, + scheduling_rank=participants[0], + ) + except Exception: + # MpiConsensusTransport owns a duplicated communicator as soon as + # construction returns. The coordinator is not yet published to + # shutdown(), so roll that ownership back transactionally while + # preserving the original constructor exception. + try: + transport.close(_CONSENSUS_STARTUP_CLOSE_TIMEOUT_S) + except Exception as close_error: + logger.error( + "Python asynchronous-consensus startup rollback failed " + f"to close its transport: {close_error}" + ) + raise + self._async_consensus = coordinator + logger.info( + "PYTHON_ASYNC_CONSENSUS transition=mode_active " + f"version={PROTOCOL_VERSION} side=ctx rank={self._dist.rank} " + f"terminal={int(terminal_requested)} " + f"peer_ready={int(peer_ready_requested)} " + f"participants={participants}" + ) + return endpoints + def _broadcast_instance_name(self) -> str: if self._dist.rank == 0: name = str(uuid.uuid4()) @@ -138,11 +432,28 @@ def _init_sync_policy(self): ) def _exchange_rank_info(self): - endpoints = cast(list, self._dist.allgather(self._transfer_worker.sender_endpoint)) + endpoint = self._transfer_worker.sender_endpoint + raw_opt_in = self._mapping.world_size > 1 and ( + self._async_terminal_flag_value != "0" or self._async_peer_ready_flag_value != "0" + ) + contribution: Any = endpoint + if raw_opt_in: + contribution = ( + _ASYNC_STARTUP_TAG, + endpoint, + self._async_startup_descriptor(), + ) + gathered = list(self._dist.allgather(contribution)) layer_num = len(self._kv_cache_manager.pp_layers) if isinstance(self._kv_cache_manager, MambaHybridCacheManager): layer_num += len(self._kv_cache_manager._impl.mamba_layer_offsets) layer_num_per_pp = cast(list, getattr(self._dist, "pp_allgather")(layer_num)) + # Validate only after every rank has entered both pre-existing startup + # exchanges. This lets a mixed old/new worker group with an accidental + # explicit opt-in fail on both sides instead of leaving the old worker + # blocked in the layer-count exchange. Flags-off workers still send + # the exact legacy endpoint value and add no collective. + endpoints = self._complete_async_consensus_startup(gathered) self._transfer_worker.populate_instance_and_rank_info( endpoints=endpoints, layer_num_per_pp=layer_num_per_pp ) @@ -151,18 +462,133 @@ def _exchange_rank_info(self): logger.info(f"self._context_info_endpoint: {self._context_info_endpoint}") def shutdown(self): - if getattr(self, "_shutdown", False): + if getattr(self, "_shutdown_complete", False): + self._log_context_activation_digest() return + worker_event = getattr(self, "_shutdown_worker_event", None) + if worker_event is not None and not worker_event.is_set(): + return worker_event self._shutdown = True - for session in list(self._send_sessions.values()): - session.close() - for session in list(self._recv_sessions.values()): - session.close() - self._send_sessions.clear() - self._send_reqs.clear() - self._recv_sessions.clear() - self._recv_reqs.clear() - self._transfer_worker.shutdown() + self._async_ready_idle_wakeup.set() + shutdown_errors: list[tuple[str, Exception]] = [] + + def run_shutdown_step(name: str, callback: Callable[[], None]) -> None: + try: + callback() + except Exception as error: + logger.error(f"Python transceiver shutdown step {name!r} failed: {error}") + shutdown_errors.append((name, error)) + + def cancel_metadata_leases() -> None: + for rid in list(self._async_ready_metadata_leases): + run_shutdown_step( + f"cancel metadata-only request {rid}", + lambda rid=rid: self._cancel_async_ready_metadata_lease(rid), + ) + self._shutdown_metadata_leases_complete = not bool(self._async_ready_metadata_leases) + + if not getattr(self, "_shutdown_metadata_leases_complete", False): + cancel_metadata_leases() + + if not getattr(self, "_shutdown_sessions_complete", False): + for rid, session in list(self._send_sessions.items()): + run_shutdown_step(f"close TxSession {rid}", session.close) + for rid, session in list(self._recv_sessions.items()): + run_shutdown_step(f"close RxSession {rid}", session.close) + self._send_sessions.clear() + self._send_reqs.clear() + self._recv_sessions.clear() + self._recv_reqs.clear() + getattr(self, "_legacy_failed_sessions", set()).clear() + # Session close is best effort. Transfer-worker shutdown is the + # final drain, so session close must not be retried after the + # worker may already have stopped. + self._shutdown_sessions_complete = True + + # A stale transceiver-side metadata lease can race native TxSession + # creation. Closing sessions above transfers ownership back to the + # metadata-only cancellation path, so retry once synchronously before + # deciding that teardown is incomplete. + if not getattr(self, "_shutdown_metadata_leases_complete", False): + cancel_metadata_leases() + if not self._shutdown_metadata_leases_complete: + shutdown_errors.append( + ( + "metadata-only requests", + RuntimeError( + "Python transceiver metadata-only request teardown is incomplete" + ), + ) + ) + + if not getattr(self, "_shutdown_consensus_complete", False): + if self._async_consensus is None: + self._shutdown_consensus_complete = True + else: + errors_before_consensus = len(shutdown_errors) + + def shutdown_consensus() -> None: + self._async_consensus.shutdown() + logger.info( + "PYTHON_ASYNC_CONSENSUS transition=shutdown_summary " + f"rank={self._dist.rank} " + f"counters={dict(self._async_consensus_counters)}" + ) + + run_shutdown_step("asynchronous consensus", shutdown_consensus) + if len(shutdown_errors) == errors_before_consensus: + self._shutdown_consensus_complete = True + + if getattr(self, "_shutdown_metadata_leases_complete", False) and not getattr( + self, "_shutdown_worker_complete", False + ): + try: + worker_result = self._transfer_worker.shutdown() + except Exception as error: + # A non-progress-thread shutdown raises only after native + # teardown is terminal, so registered memory can now be + # released even though the error still has to be surfaced. + logger.error(f"Python transceiver shutdown step 'transfer worker' failed: {error}") + shutdown_errors.append(("transfer worker", error)) + self._shutdown_worker_complete = True + self._shutdown_worker_event = None + else: + if isinstance(worker_result, threading.Event): + # Internal progress threads cannot join themselves. The + # returned event is set when deferred native teardown is + # terminal; the owner must wait and call shutdown again + # to surface any recorded native error. + self._shutdown_worker_event = worker_result + elif worker_result is None: + self._shutdown_worker_complete = True + self._shutdown_worker_event = None + else: + raise TypeError( + "TransferWorker.shutdown() must return None or threading.Event, " + f"got {type(worker_result).__name__}" + ) + + self._shutdown_complete = ( + getattr(self, "_shutdown_metadata_leases_complete", False) + and getattr(self, "_shutdown_sessions_complete", False) + and getattr(self, "_shutdown_consensus_complete", False) + and getattr(self, "_shutdown_worker_complete", False) + ) + if self._shutdown_complete: + self._log_context_activation_digest() + worker_event = getattr(self, "_shutdown_worker_event", None) + if worker_event is not None and not getattr(self, "_shutdown_worker_complete", False): + deferred_errors = getattr(self, "_shutdown_deferred_errors", []) + deferred_errors.extend(shutdown_errors) + self._shutdown_deferred_errors = deferred_errors + return worker_event + shutdown_errors = [ + *getattr(self, "_shutdown_deferred_errors", []), + *shutdown_errors, + ] + self._shutdown_deferred_errors = [] + if shutdown_errors: + raise shutdown_errors[0][1] def __enter__(self): return self @@ -365,49 +791,528 @@ def _intersection(all_lists: List[List[int]], n_ranks: int) -> set: return {rid for rid, c in cnt.items() if c == n_ranks} def _consensus_outcome( - self, to_process, cancelled, failed, completed, allgather: Callable, need_sync: bool + self, + to_process, + known_ids, + cancelled, + cancel_quiescent, + failed, + failed_quiescent, + completed, + allgather: Callable, + need_sync: bool, ): - # CANCELLED/FAILED on any rank → global; COMPLETED only when ALL ranks agree. - # Batch the three id lists into one allgather to cut the per-step collective count. + # A cancellation decision is global when ANY rank observes it, but its + # resources are reclaimable only when ALL ranks report local + # quiescence. Failure uses the same decision/acknowledgement split; + # ERROR alone does not prove that sibling native operations drained. + # COMPLETED is global only when every rank agrees. Keep all five lists + # in one packed allgather so the acknowledgements add no rendezvous. if not need_sync: - all_c, all_f, all_done = [list(cancelled)], [list(failed)], [list(completed)] + all_c = [list(cancelled)] + all_cq = [list(cancel_quiescent)] + all_f = [list(failed)] + all_fq = [list(failed_quiescent)] + all_done = [list(completed)] else: - packed = list(allgather([list(cancelled), list(failed), list(completed)])) + packed = list( + allgather( + [ + list(cancelled), + list(cancel_quiescent), + list(failed), + list(failed_quiescent), + list(completed), + ] + ) + ) all_c = [p[0] for p in packed] - all_f = [p[1] for p in packed] - all_done = [p[2] for p in packed] + all_cq = [p[1] for p in packed] + all_f = [p[2] for p in packed] + all_fq = [p[3] for p in packed] + all_done = [p[4] for p in packed] n = len(all_c) global_cancelled = self._union(all_c) + global_cancel_quiescent = self._intersection(all_cq, n) global_failed = self._union(all_f) + global_failed_quiescent = self._intersection(all_fq, n) global_completed = self._intersection(all_done, n) - new_cancelled = [rid for rid in to_process if rid in global_cancelled] + new_cancelled = [rid for rid in known_ids if rid in global_cancelled] + reclaimable_cancelled = [rid for rid in new_cancelled if rid in global_cancel_quiescent] cancel_set = set(new_cancelled) - new_failed = [rid for rid in to_process if rid in global_failed and rid not in cancel_set] + new_failed = [rid for rid in known_ids if rid in global_failed and rid not in cancel_set] + reclaimable_failed = [rid for rid in new_failed if rid in global_failed_quiescent] terminal = cancel_set | set(new_failed) new_completed = [ rid for rid in to_process if rid in global_completed and rid not in terminal ] - return new_cancelled, new_failed, new_completed + return ( + new_cancelled, + reclaimable_cancelled, + new_failed, + reclaimable_failed, + new_completed, + ) - def _gen_consensus_outcome(self, to_process, cancelled, failed, completed): + def _gen_consensus_outcome( + self, + to_process, + known_ids, + cancelled, + cancel_quiescent, + failed, + failed_quiescent, + completed, + ): return self._consensus_outcome( - to_process, cancelled, failed, completed, self._gen_allgather, self._gen_need_sync + to_process, + known_ids, + cancelled, + cancel_quiescent, + failed, + failed_quiescent, + completed, + self._gen_allgather, + self._gen_need_sync, ) - def _ctx_consensus_outcome(self, to_process, cancelled, failed, completed, timed_out): + def _ctx_consensus_outcome( + self, + to_process, + known_ids, + cancelled, + cancel_quiescent, + failed, + failed_quiescent, + completed, + timed_out, + ): # TP first, then PP. timed_out is local-only (back-off signal). - c, f, d = self._consensus_outcome( + c, cq, f, fq, d = self._consensus_outcome( to_process, + known_ids, cancelled, + cancel_quiescent, failed, + failed_quiescent, completed, self._dist.tp_allgather, self._ctx_need_tp_sync, ) if self._ctx_need_pp_sync: pp_allgather: Callable = getattr(self._dist, "pp_allgather") - c, f, d = self._consensus_outcome(to_process, c, f, d, pp_allgather, True) - return c, f, d, timed_out + c, cq, f, fq, d = self._consensus_outcome( + to_process, known_ids, c, cq, f, fq, d, pp_allgather, True + ) + return c, cq, f, fq, d, timed_out + + def _record_async_transition( + self, + transition: str, + request_id: int, + epoch: int, + outcome: Optional[ConsensusOutcome] = None, + ) -> None: + self._async_consensus_counters[transition] += 1 + count = self._async_consensus_counters[transition] + if count != 1 and count % 512 != 0: + return + outcome_text = "" if outcome is None else f" outcome={outcome.name}" + logger.info( + "PYTHON_ASYNC_CONSENSUS " + f"transition={transition} count={count} " + f"rank={self._dist.rank} request_id={request_id} epoch={epoch}" + f"{outcome_text}" + ) + + @staticmethod + def _retire_async_epoch( + epochs: OrderedDict[int, int], request_id: int, next_epoch: int + ) -> None: + """Retain a bounded replay window for recently retired request IDs.""" + epochs[request_id] = max(next_epoch, epochs.get(request_id, 0)) + epochs.move_to_end(request_id) + while len(epochs) > _MAX_RETIRED_CONSENSUS_REQUESTS: + epochs.popitem(last=False) + + def _progress_async_consensus(self) -> bool: + coordinator = self._async_consensus + if coordinator is None: + return False + events = coordinator.poll() + for event in events: + if event.kind == ConsensusEventKind.TERMINAL_COMMIT: + published_epoch = self._async_terminal_published.get(event.request_id) + if published_epoch != event.epoch: + raise RuntimeError( + "received an unexpected asynchronous terminal commit: " + f"request_id={event.request_id}, epoch={event.epoch}, " + f"published_epoch={published_epoch}" + ) + self._async_terminal_commits[event.request_id] = event + self._record_async_transition( + "terminal_commit", + event.request_id, + event.epoch, + event.outcome, + ) + elif event.kind == ConsensusEventKind.READY_PREPARE: + self._apply_async_ready_prepare(event) + elif event.kind == ConsensusEventKind.READY_RELEASE: + self._apply_async_ready_release(event) + elif event.kind == ConsensusEventKind.READY_COMPLETE: + self._apply_async_ready_complete(event) + elif event.kind == ConsensusEventKind.READY_ABORT: + self._apply_async_ready_abort(event) + elif event.kind == ConsensusEventKind.READY_ABORT_FINALIZE: + self._apply_async_ready_abort_finalize(event) + else: + raise RuntimeError(f"unhandled asynchronous consensus event: {event}") + return bool(events) + + def _apply_async_ready_prepare(self, event: ConsensusEvent) -> None: + coordinator = cast(AsyncConsensusCoordinator, self._async_consensus) + key = (event.request_id, event.epoch) + if self._async_ready_published.get(event.request_id) != event.epoch: + raise RuntimeError(f"READY_PREPARE has no matching local publication: {event}") + req = self._wait_reqs.pop(event.request_id, None) + if req is None: + raise RuntimeError(f"READY_PREPARE has no waiting request: {event}") + # PREPARE is a hidden, abortable lease. Do not expose the request to + # an ordinary scheduler yet: another rank may still withdraw, and a + # state-only rollback cannot undo KV/scheduler mutations. Rank zero + # becomes visible at READY_RELEASE; followers become visible only when + # the exact rank-zero PP schedule contains this request. + self._async_ready_prepared[key] = req + coordinator.acknowledge_ready(event.request_id, event.epoch) + self._async_ready_acknowledged.add(key) + self._record_async_transition("ready_prepare", event.request_id, event.epoch) + + def _apply_async_ready_release(self, event: ConsensusEvent) -> None: + if self._async_ready_published.get(event.request_id) != event.epoch: + raise RuntimeError(f"READY_RELEASE has no matching local publication: {event}") + key = (event.request_id, event.epoch) + req = self._async_ready_prepared.get(key) + if req is None: + raise RuntimeError(f"READY_RELEASE has no prepared request: {event}") + req.state = LlmRequestState.CONTEXT_INIT + self._async_ready_released.add(key) + self._record_async_transition("ready_release", event.request_id, event.epoch) + + def _apply_async_ready_complete(self, event: ConsensusEvent) -> None: + key = (event.request_id, event.epoch) + req = self._async_ready_activated.pop(key, None) + if req is None: + raise RuntimeError(f"READY_COMPLETE has no activated request: {event}") + self._async_ready_released.discard(key) + self._finish_async_ready_epoch(event.request_id, event.epoch) + self._record_async_transition("ready_complete", event.request_id, event.epoch) + + def _apply_async_ready_abort(self, event: ConsensusEvent) -> None: + coordinator = cast(AsyncConsensusCoordinator, self._async_consensus) + key = (event.request_id, event.epoch) + if key in self._async_ready_aborted: + coordinator.acknowledge_ready_abort(event.request_id, event.epoch) + return + req = self._async_ready_prepared.pop(key, None) + if req is None: + req = self._wait_reqs.pop(event.request_id, None) + if req is not None: + req.state = LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER + setattr(req, _ASYNC_READY_CANCELLED_EPOCH_ATTR, event.epoch) + self._async_ready_aborted[key] = req + coordinator.acknowledge_ready_abort(event.request_id, event.epoch) + self._record_async_transition("ready_abort", event.request_id, event.epoch) + + def _apply_async_ready_abort_finalize(self, event: ConsensusEvent) -> None: + key = (event.request_id, event.epoch) + if key not in self._async_ready_aborted: + raise RuntimeError(f"READY_ABORT_FINALIZE has no aborted request: {event}") + req = self._async_ready_aborted.pop(key) + if req is None: + # The protocol round is complete, but a request already in flight + # through the executor may reach this rank after the final event. + # Retain a bounded integration tombstone so that first late + # materialization is excluded rather than being published as the + # next readiness epoch. + self._async_ready_finalized_without_request[event.request_id] = event.epoch + self._async_ready_finalized_without_request.move_to_end(event.request_id) + while ( + len(self._async_ready_finalized_without_request) > _MAX_RETIRED_CONSENSUS_REQUESTS + ): + self._async_ready_finalized_without_request.popitem(last=False) + if not self._cancel_async_ready_metadata_lease(event.request_id): + raise RuntimeError( + "READY_ABORT_FINALIZE could not retire metadata-only " + f"ownership: request_id={event.request_id}, " + f"epoch={event.epoch}" + ) + self._finish_async_ready_epoch(event.request_id, event.epoch) + self._record_async_transition("ready_abort_finalize", event.request_id, event.epoch) + + def _bind_async_ready_abort(self, req: LlmRequest) -> bool: + """Bind a request that materialized after its readiness abort. + + READY_ABORT acknowledgement is a protocol action and cannot wait for + executor timing. Conversely, a late local request must not be treated + as a fresh readiness vote. Bind it to an active request-less abort or + consume the bounded post-finalize tombstone and leave it in the + cancellation wait state. + """ + request_id = get_unique_rid(req) + active_key = next( + ( + key + for key, aborted_req in self._async_ready_aborted.items() + if key[0] == request_id and aborted_req is None + ), + None, + ) + if active_key is not None: + epoch = active_key[1] + self._async_ready_aborted[active_key] = req + else: + epoch = self._async_ready_finalized_without_request.pop(request_id, None) + if epoch is None: + return False + req.state = LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER + setattr(req, _ASYNC_READY_CANCELLED_EPOCH_ATTR, epoch) + return True + + def _finish_async_ready_epoch(self, request_id: int, epoch: int) -> None: + if self._async_ready_published.get(request_id) == epoch: + del self._async_ready_published[request_id] + key = (request_id, epoch) + self._async_ready_acknowledged.discard(key) + self._async_ready_withdrawn.discard(key) + self._retire_async_epoch(self._async_ready_epoch, request_id, epoch + 1) + + @staticmethod + def _terminal_outcome_from_status( + status: SessionStatus, + ) -> Optional[ConsensusOutcome]: + if status == SessionStatus.CANCELLED: + return ConsensusOutcome.CANCELLED + if status == SessionStatus.ERROR: + return ConsensusOutcome.FAILED + if status in ( + SessionStatus.KV_TRANSFERRED, + SessionStatus.FULLY_TRANSFERRED, + ): + return ConsensusOutcome.COMPLETED + return None + + def _seal_and_snapshot_terminal( + self, + session: TxSessionBase, + result: Optional[WaitResult], + ) -> Optional[ConsensusOutcome]: + """Seal ``session`` and return its immutable local terminal vote. + + New native sessions expose an atomic terminal snapshot. Keep the + fallback until every transfer backend implements that API; it re-reads + status only after the old seal boundary so cancellation cannot replace + a previously observed success between observation and publication. + """ + snapshot_terminal = getattr(session, "seal_and_snapshot_terminal", None) + if snapshot_terminal is not None: + status = snapshot_terminal() + if status is None: + return None + return self._terminal_outcome_from_status(status) + + candidate = self._terminal_outcome_from_status(session.status) + if candidate is None: + if result == WaitResult.FAILED or session.has_failed(): + candidate = ConsensusOutcome.FAILED + elif result == WaitResult.COMPLETED or session.is_completed(): + candidate = ConsensusOutcome.COMPLETED + if candidate is None: + return None + if not session.seal_and_check_quiescent(): + return None + outcome = self._terminal_outcome_from_status(session.status) + if outcome is not None: + return outcome + return candidate + + def _publish_async_ctx_terminal_votes( + self, + block_all: bool, + request_ids: Optional[set[int]] = None, + ) -> list[int]: + coordinator = cast(AsyncConsensusCoordinator, self._async_consensus) + timed_out: list[int] = [] + for rid, session in list(self._send_sessions.items()): + if request_ids is not None and rid not in request_ids: + continue + if rid in self._async_terminal_published: + continue + if session.has_transferring_tasks(): + continue + + result = session.wait_complete(blocking=block_all) + outcome = self._seal_and_snapshot_terminal(session, result) + if outcome is None: + if result == WaitResult.TIMEOUT: + timed_out.append(rid) + logger.warning( + f"TxSession rid={session.disagg_request_id} timed out after " + f"{self._sender_future_timeout_ms}ms" + ) + continue + epoch = self._async_terminal_epoch.get(rid, 0) + coordinator.publish_terminal(rid, outcome, epoch) + self._async_terminal_published[rid] = epoch + self._record_async_transition("terminal_vote", rid, epoch, outcome) + return timed_out + + def _apply_async_ctx_terminal_commits( + self, + *, + mark_complete: bool, + only_request_id: Optional[int] = None, + request_ids: Optional[set[int]] = None, + ) -> tuple[list[int], list[int]]: + completed: list[int] = [] + failed: list[int] = [] + for rid, event in list(self._async_terminal_commits.items()): + if only_request_id is not None and rid != only_request_id: + continue + if request_ids is not None and rid not in request_ids: + continue + session = self._send_sessions.get(rid) + req = self._send_reqs.get(rid) + if session is None or req is None: + raise RuntimeError( + f"terminal commit arrived after local resources were released: rid={rid}" + ) + if session.has_transferring_tasks(): + raise RuntimeError(f"terminal commit arrived before local quiescence: rid={rid}") + + if event.outcome == ConsensusOutcome.COMPLETED: + if mark_complete: + req.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE + completed.append(rid) + elif event.outcome == ConsensusOutcome.FAILED: + req.state = LlmRequestState.DISAGG_TRANS_ERROR + failed.append(rid) + elif event.outcome == ConsensusOutcome.CANCELLED: + # The ordinary status poll can observe the authoritative + # cancellation before the executor retries its cancellation + # request. Preserve a request-scoped tombstone so the + # AsyncTransferManager can acknowledge and retire exactly its + # transceiver leg instead of losing ownership permanently. + self._async_terminal_cancelled[rid] = (event.epoch, req) + self._record_context_cancelled_request_id(rid) + else: + raise RuntimeError(f"invalid terminal commit outcome: {event}") + + session.close() + del self._send_reqs[rid] + del self._send_sessions[rid] + del self._async_terminal_commits[rid] + self._async_terminal_published.pop(rid, None) + if event.outcome != ConsensusOutcome.CANCELLED: + self._retire_async_epoch(self._async_terminal_epoch, rid, event.epoch + 1) + return completed, failed + + def _record_context_cancelled_request_id(self, request_id: int) -> None: + request_ids = getattr(self, "_context_cancelled_request_ids", None) + if request_ids is None: + request_ids = [] + self._context_cancelled_request_ids = request_ids + request_ids.append(request_id) + + def take_context_cancelled_request_ids(self) -> List[int]: + """Transfer newly quiescent CTX cancellation IDs to the executor.""" + pending = getattr(self, "_context_cancelled_request_ids", None) + if pending is None: + return [] + request_ids = list(pending) + pending.clear() + return request_ids + + def _acknowledge_async_terminal_cancellation(self, request_id: int, req: LlmRequest) -> bool: + cancelled_by_request = getattr(self, "_async_terminal_cancelled", None) + if cancelled_by_request is None: + return False + cancelled = cancelled_by_request.get(request_id) + if cancelled is None: + return False + epoch, owned_req = cancelled + if owned_req is not req: + raise RuntimeError( + "request ID was reused before an asynchronous terminal " + f"cancellation was acknowledged: request_id={request_id}, epoch={epoch}" + ) + del cancelled_by_request[request_id] + # A local user-cancellation retry can consume the commit before the + # executor's ordinary status poll drains this handoff queue. Remove + # that now-local acknowledgement so it cannot later be misclassified + # as a peer cancellation after the request has already terminated. + pending = getattr(self, "_context_cancelled_request_ids", None) + if pending: + pending[:] = [pending_id for pending_id in pending if pending_id != request_id] + self._retire_async_epoch(self._async_terminal_epoch, request_id, epoch + 1) + self._record_async_transition("terminal_cancel_ack", request_id, epoch) + return True + + def _check_context_transfer_status_async( + self, at_least_request_num: Optional[int], mark_complete: bool + ) -> tuple[list[int], list[int]]: + snapshot_ids = set(self._send_sessions) + drain_timeout_ms: Optional[int] = None + if at_least_request_num is None: + target_commits = len(snapshot_ids) + drain_timeout_ms = ( + self._sender_future_timeout_ms or getattr(self, "kv_transfer_timeout_ms", 0) or 0 + ) + timeout_s = max(0, drain_timeout_ms) / 1000.0 + else: + target_commits = min(max(0, at_least_request_num), len(snapshot_ids)) + timeout_s = max(0, self.kv_transfer_poll_interval_ms or 0) / 1000.0 + deadline = time.monotonic() + timeout_s + completed: list[int] = [] + failed: list[int] = [] + committed = 0 + + while True: + self._progress_async_consensus() + # Poll non-blockingly and use a single deadline for the fixed + # entry snapshot. This preserves the configured timeout without + # multiplying it by the number of sessions. + self._publish_async_ctx_terminal_votes( + block_all=False, + request_ids=snapshot_ids, + ) + self._progress_async_consensus() + committed_ids = snapshot_ids.intersection(self._async_terminal_commits) + new_completed, new_failed = self._apply_async_ctx_terminal_commits( + mark_complete=mark_complete, + request_ids=snapshot_ids, + ) + completed.extend(new_completed) + failed.extend(new_failed) + committed += len(committed_ids) + + if at_least_request_num == 0 or committed >= target_commits: + break + remaining_s = deadline - time.monotonic() + if remaining_s <= 0: + break + time.sleep(min(0.001, remaining_s)) + + if at_least_request_num is None: + remaining_ids = snapshot_ids.intersection(self._send_sessions) + if remaining_ids: + logger.warning( + "Timed out draining asynchronous context transfer snapshot " + f"after {drain_timeout_ms}ms: " + f"remaining_rids={sorted(remaining_ids)}" + ) + self._transfer_worker.sweep_stale_req_infos() + return completed, failed def _sync_transfer_timing(self, reqs: list): """Allgather timing for a batch of completed requests in one collective. @@ -512,8 +1417,18 @@ def _apply_aux(self, session, req: LlmRequest): def _get_or_create_send_session(self, req: LlmRequest) -> TxSessionBase: rid = get_unique_rid(req) assert rid is not None + if rid in getattr(self, "_async_terminal_cancelled", {}): + raise RuntimeError( + "cannot reuse a request ID before its asynchronous terminal " + f"cancellation is acknowledged: request_id={rid}" + ) if rid not in self._send_sessions: self._send_sessions[rid] = self._transfer_worker.create_tx_session(req) + self._async_ready_metadata_leases.discard(rid) + # Publish request ownership before any native dispatch. A pre-cancelled + # session, or a cancellation racing send/send_aux, must never leave a + # session without its matching request bookkeeping. + self._send_reqs[rid] = req return self._send_sessions[rid] def _finalize_send(self, req: LlmRequest, session: TxSessionBase): @@ -531,7 +1446,6 @@ def _finalize_send(self, req: LlmRequest, session: TxSessionBase): ctx_dp_rank=self._dp_rank, disagg_info_endpoint=self._context_info_endpoint, ) - self._send_reqs[rid] = req @nvtx_range("KvCacheTransceiverV2.respond_and_send_async") def respond_and_send_async(self, req: LlmRequest): @@ -539,8 +1453,20 @@ def respond_and_send_async(self, req: LlmRequest): req.set_kv_cache_transfer_start(tensorrt_llm.bindings.global_steady_clock_now()) session = self._get_or_create_send_session(req) req.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS - session.send(self._create_kv_slice(req)) - self._finalize_send(req, session) + if session.status == SessionStatus.CANCELLED: + return + kv_slice = self._create_kv_slice(req) + try: + session.send(kv_slice) + if session.status == SessionStatus.CANCELLED: + return + self._finalize_send(req, session) + except RuntimeError: + # Native pre-cancellation seals the session. Retain the paired + # maps so the ordinary cancellation/status path owns cleanup. + if session.status == SessionStatus.CANCELLED: + return + raise @nvtx_range("KvCacheTransceiverV2.request_and_receive_sync") def request_and_receive_sync(self, req: LlmRequest): @@ -591,14 +1517,25 @@ def request_and_receive_async(self, req: LlmRequest): req.state = LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS session = self._transfer_worker.create_rx_session(req) self._recv_sessions[rid] = session + self._recv_reqs[rid] = req + if session.status == SessionStatus.CANCELLED: + return kv_slice = self._create_kv_slice(req) req.py_kv_cache_xfer_bytes = self._slice_num_bytes(kv_slice) * self._kv_size_rank_factor - session.receive(kv_slice) - self._recv_reqs[rid] = req + try: + session.receive(kv_slice) + except RuntimeError: + if session.status == SessionStatus.CANCELLED: + return + raise def check_context_transfer_status( self, at_least_request_num: Optional[int], mark_complete: bool = False ): + if getattr(self, "_async_terminal_consensus_enabled", False): + return self._check_context_transfer_status_async(at_least_request_num, mark_complete) + if getattr(self, "_async_consensus", None) is not None: + self._progress_async_consensus() # A worker that never sends KV has nothing to reconcile here, so skip the consensus. Safe # because the flag flips together on every rank and never resets, so they all skip in step; # gating on the live session dict instead would not be, since a cancel clears it per-rank. @@ -618,12 +1555,30 @@ def check_context_transfer_status( block_all, ) - completed, timed_out, failed, cancelled = [], [], [], [] + known_ids = list(self._send_sessions) + completed, timed_out, failed, failed_quiescent = [], [], [], [] + cancelled, cancel_quiescent = [], [] + for rid, session in self._send_sessions.items(): + if session.status == SessionStatus.CANCELLED: + cancelled.append(rid) + if not session.has_transferring_tasks(): + cancel_quiescent.append(rid) + elif rid in self._legacy_failed_sessions or session.has_failed(): + failed.append(rid) + if session.seal_and_check_quiescent(): + failed_quiescent.append(rid) for rid in to_process: session = self._send_sessions[rid] result = session.wait_complete(blocking=block_all) if session.status == SessionStatus.CANCELLED: - cancelled.append(rid) + if rid in failed: + failed.remove(rid) + if rid in failed_quiescent: + failed_quiescent.remove(rid) + if rid not in cancelled: + cancelled.append(rid) + if rid not in cancel_quiescent and not session.has_transferring_tasks(): + cancel_quiescent.append(rid) elif result == WaitResult.COMPLETED: completed.append(rid) elif result is None: @@ -635,17 +1590,44 @@ def check_context_transfer_status( timed_out.append(rid) else: logger.warning(f"TxSession rid={session.disagg_request_id} failed") - failed.append(rid) + if rid not in failed: + failed.append(rid) + if rid not in failed_quiescent and session.seal_and_check_quiescent(): + failed_quiescent.append(rid) # All ranks must agree on per-rid outcome to avoid req.state divergence. - cancelled, failed, completed, timed_out = self._ctx_consensus_outcome( - to_process, cancelled, failed, completed, timed_out + cancelled, reclaimable_cancelled, failed, reclaimable_failed, completed, timed_out = ( + self._ctx_consensus_outcome( + to_process, + known_ids, + cancelled, + cancel_quiescent, + failed, + failed_quiescent, + completed, + timed_out, + ) ) for rid in cancelled: + self._legacy_failed_sessions.discard(rid) + session = self._send_sessions[rid] + # A peer may be the first rank to observe cancellation. Apply + # that decision locally, but retain the request and session while + # a native write is still active. CANCELLED is a decision, not a + # reclamation acknowledgement. + session.cancel() + + for rid in failed: + self._legacy_failed_sessions.add(rid) + self._send_sessions[rid].seal_and_check_quiescent() + + for rid in reclaimable_cancelled: + self._legacy_failed_sessions.discard(rid) self._send_sessions[rid].close() del self._send_reqs[rid] del self._send_sessions[rid] + self._record_context_cancelled_request_id(rid) for rid in completed: if mark_complete: @@ -653,15 +1635,19 @@ def check_context_transfer_status( self._send_sessions[rid].close() del self._send_reqs[rid] del self._send_sessions[rid] - self._close_failed_sessions(self._send_sessions, self._send_reqs, failed) + for rid in reclaimable_failed: + self._legacy_failed_sessions.discard(rid) + self._close_failed_sessions(self._send_sessions, self._send_reqs, reclaimable_failed) # Sweep orphaned RecvReqInfo entries from ADP broadcast on non-assigned # DP ranks (entries that will never have a TxSession created for them). self._transfer_worker.sweep_stale_req_infos() - return completed, failed + return completed, reclaimable_failed def check_gen_transfer_status(self, at_least_request_num: Optional[int]): + if getattr(self, "_async_consensus", None) is not None: + self._progress_async_consensus() if not self._ever_had_recv_session and not self._gen_need_sync: return [], [], [] block_all = at_least_request_num is None @@ -678,7 +1664,18 @@ def check_gen_transfer_status(self, at_least_request_num: Optional[int]): block_all, ) - completed, failed, cancelled = [], [], [] + known_ids = list(self._recv_sessions) + completed, failed, failed_quiescent = [], [], [] + cancelled, cancel_quiescent = [], [] + for rid, session in self._recv_sessions.items(): + if session.status == SessionStatus.CANCELLED: + cancelled.append(rid) + if not session.has_transferring_tasks(): + cancel_quiescent.append(rid) + elif rid in self._legacy_failed_sessions or session.has_failed(): + failed.append(rid) + if session.seal_and_check_quiescent(): + failed_quiescent.append(rid) for rid in to_process: session = self._recv_sessions[rid] result = session.wait_complete(blocking=block_all) @@ -687,7 +1684,14 @@ def check_gen_transfer_status(self, at_least_request_num: Optional[int]): # cancel) or by a remote CANCEL_SESSION message (e.g. CTX # server timeout). Return the req objects so the caller can # distinguish the two cases and set the appropriate state. - cancelled.append(rid) + if rid in failed: + failed.remove(rid) + if rid in failed_quiescent: + failed_quiescent.remove(rid) + if rid not in cancelled: + cancelled.append(rid) + if rid not in cancel_quiescent and not session.has_transferring_tasks(): + cancel_quiescent.append(rid) elif result == WaitResult.COMPLETED: req = self._recv_reqs[rid] if session.transfer_end_time is not None: @@ -696,18 +1700,40 @@ def check_gen_transfer_status(self, at_least_request_num: Optional[int]): req.set_kv_cache_size(session.kv_cache_size_bytes) completed.append(rid) elif result == WaitResult.FAILED: - failed.append(rid) + if rid not in failed: + failed.append(rid) + if rid not in failed_quiescent and session.seal_and_check_quiescent(): + failed_quiescent.append(rid) # else: None — KV done but aux still in flight; re-poll next cycle # All ranks must agree on per-rid outcome to avoid req.state divergence. - cancelled, failed, completed = self._gen_consensus_outcome( - to_process, cancelled, failed, completed + cancelled, reclaimable_cancelled, failed, reclaimable_failed, completed = ( + self._gen_consensus_outcome( + to_process, + known_ids, + cancelled, + cancel_quiescent, + failed, + failed_quiescent, + completed, + ) ) cancelled_reqs = [] for rid in cancelled: + self._legacy_failed_sessions.discard(rid) + session = self._recv_sessions[rid] + # Native receive cancellation is not reclaimable until both + # active writes and sender acknowledgements have drained. + session.cancel() + for rid in failed: + self._legacy_failed_sessions.add(rid) + self._recv_sessions[rid].seal_and_check_quiescent() + for rid in reclaimable_cancelled: + self._legacy_failed_sessions.discard(rid) + session = self._recv_sessions[rid] cancelled_reqs.append(self._recv_reqs[rid]) - self._recv_sessions[rid].close() + session.close() del self._recv_reqs[rid] del self._recv_sessions[rid] @@ -737,14 +1763,16 @@ def check_gen_transfer_status(self, at_least_request_num: Optional[int]): session.close() del self._recv_reqs[rid] del self._recv_sessions[rid] - if failed: + if reclaimable_failed: logger.warning( f"Disagg gen transfer FAILED rank={self._dist.rank} " - f"rids={failed} gen_need_sync={self._gen_need_sync}" + f"rids={reclaimable_failed} gen_need_sync={self._gen_need_sync}" ) - self._close_failed_sessions(self._recv_sessions, self._recv_reqs, failed) + for rid in reclaimable_failed: + self._legacy_failed_sessions.discard(rid) + self._close_failed_sessions(self._recv_sessions, self._recv_reqs, reclaimable_failed) - return completed, failed, cancelled_reqs + return completed, reclaimable_failed, cancelled_reqs def _poll_gen_sessions_for_poll_interval(self, wait_num: int) -> None: poll_interval_s = (self.kv_transfer_poll_interval_ms or 0) / 1000.0 @@ -797,6 +1825,23 @@ def _assert_disagg_history_declared(self, req: LlmRequest) -> None: f"_try_schedule_disagg_gen_init." ) + def owns_request(self, req: LlmRequest) -> bool: + """Return whether transfer or readiness state still owns ``req``.""" + rid = get_unique_rid(req) + return ( + rid in self._wait_reqs + or rid in self._send_sessions + or rid in self._send_reqs + or rid in self._recv_sessions + or rid in self._recv_reqs + or rid in getattr(self, "_async_terminal_cancelled", {}) + or rid in self._async_ready_published + or any(key[0] == rid for key in self._async_ready_prepared) + or any(key[0] == rid for key in self._async_ready_activated) + or any(key[0] == rid for key in self._async_ready_aborted) + or rid in self._async_ready_metadata_leases + ) + def cancel_request(self, req: LlmRequest) -> bool: """Cancel the transfer for the given request. @@ -804,9 +1849,103 @@ def cancel_request(self, req: LlmRequest) -> bool: retry next iteration. Returns True when safe to free KV memory. """ rid = get_unique_rid(req) - - # Not yet started (generation-first wait queue). - self._wait_reqs.pop(rid, None) + if getattr(self, "_async_consensus", None) is not None: + self._progress_async_consensus() + + if self._acknowledge_async_terminal_cancellation(rid, req): + return True + + # A published terminal vote is immutable. In particular, a late local + # cancellation must not mutate a session that already voted COMPLETE + # or FAILED, nor send a contradictory remote cancellation. Retain all + # resources until the coordinator's authoritative commit is applied. + if ( + getattr(self, "_async_terminal_consensus_enabled", False) + and rid in self._async_terminal_published + ): + self._apply_async_ctx_terminal_commits( + mark_complete=False, + only_request_id=rid, + ) + if self._acknowledge_async_terminal_cancellation(rid, req): + return True + return not self.owns_request(req) + + # A generation-first readiness round owns the request until its + # authoritative schedule activation completes or its abort finalizes. + # Never remove that ownership directly from the cancellation path. + waiting_req = self._wait_reqs.get(rid) + prepared_key = next((key for key in self._async_ready_prepared if key[0] == rid), None) + prepared_req = ( + self._async_ready_prepared.get(prepared_key) if prepared_key is not None else None + ) + activated_key = next( + (key for key in self._async_ready_activated if key[0] == rid), + None, + ) + activated_req = ( + self._async_ready_activated.get(activated_key) if activated_key is not None else None + ) + published_epoch = self._async_ready_published.get(rid) + readiness_owned = ( + waiting_req is not None + or prepared_req is not None + or activated_req is not None + or published_epoch is not None + or any(key[0] == rid for key in self._async_ready_aborted) + ) + if getattr(self, "_async_peer_ready_consensus_enabled", False) and readiness_owned: + self._async_ready_idle_wakeup.set() + cancelled_req = ( + waiting_req + if waiting_req is not None + else prepared_req + if prepared_req is not None + else activated_req + if activated_req is not None + else req + ) + setattr( + cancelled_req, + _ASYNC_READY_CANCELLED_EPOCH_ATTR, + published_epoch + if published_epoch is not None + else self._async_ready_epoch.get(rid, 0), + ) + if published_epoch is None: + # A peer may already have voted READY even though local peer + # metadata has not arrived. Join the same epoch with a + # withdrawal and retain the waiting request until every rank + # applies READY_ABORT and the coordinator finalizes it. + published_epoch = self._async_ready_epoch.get(rid, 0) + self._async_ready_published[rid] = published_epoch + key = (rid, published_epoch) + if key in self._async_ready_aborted: + return False + if key in self._async_ready_acknowledged: + # ACK is an irrevocable lease. Let the authoritative PP + # schedule activate it and wait for READY_COMPLETE before a + # cancellation can reclaim request resources. + return False + coordinator = cast(AsyncConsensusCoordinator, self._async_consensus) + if key not in self._async_ready_withdrawn: + if coordinator.withdraw_ready(rid, published_epoch): + self._async_ready_withdrawn.add(key) + self._record_async_transition( + "ready_withdraw", + rid, + published_epoch, + ConsensusOutcome.WITHDRAWN, + ) + else: + # The coordinator already considers the local lease + # irrevocable; retain ownership until its final event. + self._async_ready_acknowledged.add(key) + self._progress_async_consensus() + if rid in self._async_ready_published: + return False + else: + self._wait_reqs.pop(rid, None) has_transferring = False @@ -814,6 +1953,12 @@ def cancel_request(self, req: LlmRequest) -> bool: self._send_sessions[rid].cancel() if self._send_sessions[rid].has_transferring_tasks(): has_transferring = True + elif getattr(self, "_async_terminal_consensus_enabled", False): + self._publish_async_ctx_terminal_votes(block_all=False) + self._progress_async_consensus() + self._apply_async_ctx_terminal_commits(mark_complete=False, only_request_id=rid) + if rid in self._send_sessions: + has_transferring = True else: self._send_sessions[rid].close() del self._send_reqs[rid] @@ -830,7 +1975,7 @@ def cancel_request(self, req: LlmRequest) -> bool: if has_transferring: return False - return True + return self._cancel_async_ready_metadata_lease(rid) def get_disaggregated_params(self) -> Dict[str, Any]: # Keep this aligned with fields populated in respond_and_send_async(). @@ -843,23 +1988,200 @@ def get_disaggregated_params(self) -> Dict[str, Any]: # broadcasts REQUEST_DATA to all ctx DP ranks. The actual ctx_dp_rank # is stamped into ContextPhaseParams by respond_and_send_async() after # the prefill is scheduled. + if not self._publish_disaggregated_params: + return {} ctx_dp_rank = None if self._mapping.enable_attention_dp else self._dp_rank return { "ctx_dp_rank": ctx_dp_rank, "ctx_info_endpoint": [self._context_info_endpoint] if self._context_info_endpoint else None, + # The instance name is a per-transceiver UUID broadcast across the + # worker group. The generation-first coordinator uses it only to + # reject unowned/estimation-lifetime endpoint metadata; it is not + # forwarded in an inference request. + **( + {"ctx_endpoint_generation": self._instance_name} + if self._async_peer_ready_consensus_enabled + else {} + ), } + def exclude_context_requests_from_readiness(self, requests: List[LlmRequest]) -> None: + """Withdraw known-cancelled requests before readiness progression.""" + for req in requests: + rid = get_unique_rid(req) + if ( + getattr(self, "_async_peer_ready_consensus_enabled", False) + and rid not in self._send_sessions + and rid not in self._async_ready_metadata_leases + ): + # A request and its cancellation can materialize in the same + # executor fetch. Establish native metadata ownership without + # publishing READY so already-arrived or delayed REQUEST_DATA + # is retired by the normal metadata-only cancellation path. + self._transfer_worker.pin_peer_req_infos_for_send(rid) + self._async_ready_metadata_leases.add(rid) + epoch = self._async_ready_published.get(rid, self._async_ready_epoch.get(rid, 0)) + setattr(req, _ASYNC_READY_CANCELLED_EPOCH_ATTR, epoch) + + if not getattr(self, "_async_peer_ready_consensus_enabled", False): + # Terminal-only/default-off readiness still uses the legacy + # wait map. Remove an existing waiter before the immediately + # following prepare_context_requests([]) can promote it. + self._wait_reqs.pop(rid, None) + req.state = LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER + continue + prepared_key = next((key for key in self._async_ready_prepared if key[0] == rid), None) + activated_key = next( + (key for key in self._async_ready_activated if key[0] == rid), + None, + ) + readiness_owned = ( + rid in self._wait_reqs + or rid in self._async_ready_published + or prepared_key is not None + or activated_key is not None + or any(key[0] == rid for key in self._async_ready_aborted) + ) + if not readiness_owned: + # A newly fetched request has not entered a readiness round on + # any rank yet. The cancellation marker keeps it excluded when + # prepare_context_requests() runs immediately afterwards. + req.state = LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER + continue + + if rid not in self._async_ready_published: + self._async_ready_published[rid] = epoch + key = (rid, epoch) + if key in self._async_ready_acknowledged: + # The prepared lease is already irrevocable. Let its release + # finish without rolling local state back; cancellation will + # retry against the next lifecycle phase. + continue + req.state = LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER + if key in self._async_ready_aborted: + continue + if key in self._async_ready_withdrawn: + continue + coordinator = cast(AsyncConsensusCoordinator, self._async_consensus) + if coordinator.withdraw_ready(rid, epoch): + self._async_ready_withdrawn.add(key) + self._record_async_transition( + "ready_withdraw", + rid, + epoch, + ConsensusOutcome.WITHDRAWN, + ) + else: + self._async_ready_acknowledged.add(key) + + def activate_context_requests_for_schedule(self, requests: List[LlmRequest]) -> None: + """Activate readiness leases selected by the authoritative PP schedule. + + Rank zero receives ``READY_RELEASE`` and runs the only authoritative + scheduling decision. Followers call this hook after deserializing + that exact schedule and immediately before their mirrored scheduler + pass. Consequently no follower can mutate scheduler or KV state for + a request that rank zero did not select. + """ + if not getattr(self, "_async_peer_ready_consensus_enabled", False): + return + coordinator = cast(AsyncConsensusCoordinator, self._async_consensus) + activations = [] + for req in requests: + rid = get_unique_rid(req) + epoch = self._async_ready_published.get(rid) + if epoch is None: + continue + key = (rid, epoch) + prepared_req = self._async_ready_prepared.get(key) + if prepared_req is None: + # The hook receives all scheduled requests, most of which do + # not belong to a readiness round. A published-but-not-yet- + # prepared request, however, must not appear in the schedule. + if key in self._async_ready_activated: + continue + raise RuntimeError( + "authoritative PP schedule selected readiness before " + f"PREPARE: request_id={rid}, epoch={epoch}" + ) + if self._dist.rank == 0 and key not in self._async_ready_released: + raise RuntimeError( + "scheduling rank selected readiness before READY_RELEASE: " + f"request_id={rid}, epoch={epoch}" + ) + if prepared_req is not req: + raise RuntimeError( + "readiness activation request object does not match " + f"PREPARE ownership: request_id={rid}, epoch={epoch}" + ) + activations.append((rid, epoch, key, prepared_req)) + for rid, epoch, key, prepared_req in activations: + prepared_req.state = LlmRequestState.CONTEXT_INIT + del self._async_ready_prepared[key] + self._async_ready_activated[key] = prepared_req + coordinator.acknowledge_ready_activation(rid, epoch) + self._record_async_transition("ready_activate", rid, epoch) + self._record_context_activation_ids([rid for rid, _, _, _ in activations]) + + def supports_pre_active_context_requests(self) -> bool: + return bool(getattr(self, "_async_peer_ready_consensus_enabled", False)) + + def is_context_request_ready_for_activation(self, request: LlmRequest) -> bool: + if not self.supports_pre_active_context_requests(): + return False + rid = get_unique_rid(request) + epoch = self._async_ready_published.get(rid) + if epoch is None: + return False + key = (rid, epoch) + if self._async_ready_prepared.get(key) is not request: + return False + # READY_RELEASE is delivered only to the scheduling rank. Followers + # receive the authoritative activation IDs on the existing PP + # schedule, but must still prove that the same request and epoch were + # prepared locally before materializing those IDs. + return self._dist.rank != 0 or key in self._async_ready_released + + def _cancel_async_ready_metadata_lease(self, rid: int) -> bool: + if rid not in self._async_ready_metadata_leases: + return True + cancelled = self._transfer_worker.cancel_peer_req_infos_for_send(rid) + if cancelled: + self._async_ready_metadata_leases.discard(rid) + return cancelled + def prepare_context_requests(self, requests: List[LlmRequest]): # Place new generation-first context requests into wait state, then # use allgather consensus to promote ready requests to CONTEXT_INIT. for req in requests: rid = get_unique_rid(req) if rid not in self._send_sessions: + if getattr(self, "_async_peer_ready_consensus_enabled", False): + if self._bind_async_ready_abort(req): + self._wait_reqs.pop(rid, None) + continue + if getattr(req, _ASYNC_READY_CANCELLED_EPOCH_ATTR, None) is not None: + self._wait_reqs.pop(rid, None) + continue + if getattr(self, "_async_peer_ready_consensus_enabled", False): + self._transfer_worker.pin_peer_req_infos_for_send(rid) + self._async_ready_metadata_leases.add(rid) self._wait_reqs[rid] = req req.state = LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER + # Materialize this iteration's requests before polling. A READY_ABORT + # can legally arrive before this rank has voted, and its handler must + # be able to bind the request instead of fail-stopping or allowing a + # contradictory READY publication below. + if getattr(self, "_async_consensus", None) is not None: + self._progress_async_consensus() + + if getattr(self, "_async_peer_ready_consensus_enabled", False): + self._prepare_context_requests_async() + return + # Nothing waiting on any rank, so skip the consensus. The waiting set is the same on every # rank, so they all skip together. if not self._wait_reqs: @@ -874,9 +2196,46 @@ def prepare_context_requests(self, requests: List[LlmRequest]): for rid in self._wait_reqs if self._transfer_worker.has_all_peer_req_infos_for_send(rid) ] - for rid in self._ctx_consensus(local_ready): + ready_request_ids = self._ctx_consensus(local_ready) + for rid in ready_request_ids: self._wait_reqs[rid].state = LlmRequestState.CONTEXT_INIT del self._wait_reqs[rid] + self._record_context_activation_ids(ready_request_ids) + + def _prepare_context_requests_async(self) -> None: + coordinator = cast(AsyncConsensusCoordinator, self._async_consensus) + published = False + for rid in list(self._wait_reqs): + if rid in self._async_ready_published: + continue + if not self._transfer_worker.has_all_peer_req_infos_for_send(rid): + continue + epoch = self._async_ready_epoch.get(rid, 0) + coordinator.publish_ready(rid, epoch) + self._async_ready_published[rid] = epoch + self._record_async_transition("ready_vote", rid, epoch) + published = True + # Poll again so the PP-last coordinator can begin prepare immediately + # after recording its local vote. + progressed = self._progress_async_consensus() + if published or progressed or not self._wait_reqs or self._shutdown: + return + + # Native REQUEST_DATA ingress is asynchronous. A missing/stale peer + # endpoint otherwise turns the executor loop into a tight poll that can + # issue millions of legacy status collectives while making no request + # progress. Yield for one bounded configured poll interval; peer-info + # completeness remains the sole READY predicate above. Cancellation + # and shutdown are re-observed on the next iteration within this cap. + sleep_s = min( + max(0.0, self.kv_transfer_poll_interval_ms / 1000.0), + _ASYNC_READY_MAX_IDLE_SLEEP_S, + ) + if sleep_s: + self._async_ready_idle_wakeup.clear() + if self._shutdown or not self._wait_reqs: + return + self._async_ready_idle_wakeup.wait(sleep_s) def _check_compatible(self): if self._mapping.cp_size != 1: diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 0149ba7ae03a..9b98bd3664b0 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2211,6 +2211,7 @@ def create_py_executor_instance( peft_cache_config: Optional[PeftCacheConfig] = None, scheduler_config: Optional[SchedulerConfig] = None, cache_transceiver_config: Optional[CacheTransceiverConfig] = None, + publish_disaggregated_params: bool = True, virtual_memory_pools: Optional[dict] = None, execution_stream: Optional[torch.cuda.Stream] = None, dwdp_manager: Optional[DwdpManager] = None, @@ -2537,7 +2538,8 @@ def create_py_executor_instance( kv_cache_transceiver = create_kv_cache_transceiver( mapping, dist, kv_cache_manager, attention_type, - cache_transceiver_config, mamba_cache_manager) + cache_transceiver_config, mamba_cache_manager, + publish_disaggregated_params) waiting_queue_policy = (scheduler_config.waiting_queue_policy if scheduler_config is not None else diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index a686adc97528..9c1cc0b9609d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -116,7 +116,8 @@ def create_kv_cache_transceiver( kv_cache_manager: KVCacheManager, attention_type: AttentionTypeCpp, cache_transceiver_config: CacheTransceiverConfig, - mamba_cache_manager: Optional[BaseMambaCacheManager] = None): + mamba_cache_manager: Optional[BaseMambaCacheManager] = None, + publish_disaggregated_params: bool = True): if cache_transceiver_config is None or cache_transceiver_config.backend is None: logger.info("cache_transceiver is disabled") return None @@ -172,7 +173,8 @@ def create_kv_cache_transceiver( KvCacheTransceiverV2 logger.info("Using KvCacheTransceiverV2") return KvCacheTransceiverV2(mapping, dist, kv_cache_manager, - cache_transceiver_config) + cache_transceiver_config, + publish_disaggregated_params) # Default: use C++ transceiver (transceiver_runtime is None or "CPP") return BindKvCacheTransceiver(mapping, dist, kv_cache_manager, @@ -210,6 +212,45 @@ def check_gen_transfer_complete(self): def cancel_request(self, req: LlmRequest): raise NotImplementedError + def owns_request(self, req: LlmRequest) -> bool: + """Return whether the transceiver still owns request-scoped state. + + Implementations should include requests that are waiting or prepared, + not only requests with an active data transfer. The executor uses this + signal to route cancellation through the owning transceiver before it + frees request resources. + """ + return False + + def activate_context_requests_for_schedule( + self, requests: List[LlmRequest]) -> None: + """Activate generation-first requests selected by the PP schedule. + + The Python transceiver's asynchronous readiness protocol overrides + this hook. Other transceivers have no separate readiness lease, so the + authoritative schedule requires no additional action. + """ + return + + def supports_pre_active_context_requests(self) -> bool: + """Return whether generation-first CTX requests may wait outside active capacity.""" + return False + + def is_context_request_ready_for_activation(self, + request: LlmRequest) -> bool: + """Return whether rank zero may promote a pre-active CTX request.""" + return False + + def take_context_cancelled_request_ids(self) -> List[int]: + """Return newly quiescent CTX cancellations owned by the executor. + + Implementations that coordinate cancellation across ranks use this + handoff to distinguish peer-initiated cancellation from a local user + cancellation. The default implementation preserves the existing + two-list context-status API for transceivers without that signal. + """ + return [] + def supports_inflight_request_cancellation(self) -> bool: return False diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index feb02319a3e3..847d6c0972f0 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -379,6 +379,14 @@ class BatchStatePP(BatchState): microbatch_id: int = -1 +class AsyncTransferProvider(StrEnum): + """Owner of an asynchronous KV-cache transfer leg.""" + + GENERIC = "generic" + TRANSCEIVER = "transceiver" + CONNECTOR = "connector" + + class AsyncTransferManager: """ Handle asynchronous transfer of KV cache after a request has completed. @@ -395,19 +403,38 @@ class RequestTransferMetadata: def __init__(self, block_id: Optional[int]): self.block_id = block_id - self.counter = 0 + self.provider_counts: Dict[AsyncTransferProvider, int] = {} + + @property + def counter(self) -> int: + """Total leg count retained for backward-compatible inspection.""" - def start_transfer(self): - self.counter += 1 + return sum(self.provider_counts.values()) - def end_transfer(self) -> bool: + def start_transfer(self, provider: AsyncTransferProvider): + self.provider_counts[provider] = ( + self.provider_counts.get(provider, 0) + 1) + + def end_transfer(self, + provider: AsyncTransferProvider) -> Optional[bool]: """ Returns: - bool: True if there are no more transfers for this request + Optional[bool]: True if there are no more transfers for this + request, False if another leg remains, or None if the + requested provider did not own a leg. """ - self.counter -= 1 + provider_count = self.provider_counts.get(provider, 0) + if provider_count == 0: + return None + if provider_count == 1: + self.provider_counts.pop(provider) + else: + self.provider_counts[provider] = provider_count - 1 return self.counter == 0 + def has_transfer(self, provider: AsyncTransferProvider) -> bool: + return self.provider_counts.get(provider, 0) > 0 + def __init__(self, resource_manager: "ResourceManager", should_store_blocks: bool = True): @@ -427,7 +454,11 @@ def __init__(self, def requests_in_transfer(self) -> Dict[int, LlmRequest]: return self._requests_in_transfer - def start_transfer(self, request: LlmRequest): + def start_transfer( + self, + request: LlmRequest, + provider: AsyncTransferProvider = AsyncTransferProvider.GENERIC, + ): """ Called when a Cache transceiver or connector transfer is started. 1. Increment the counter for the request. @@ -435,6 +466,7 @@ def start_transfer(self, request: LlmRequest): 3. Store KV cache blocks for reuse. """ + provider = AsyncTransferProvider(provider) req_id = request.py_request_id if req_id not in self._requests_in_transfer: @@ -458,9 +490,13 @@ def start_transfer(self, request: LlmRequest): self._request_transfer_metadata[ req_id] = self.RequestTransferMetadata(block_id) - self._request_transfer_metadata[req_id].start_transfer() + self._request_transfer_metadata[req_id].start_transfer(provider) - def end_transfer(self, request: LlmRequest) -> bool: + def end_transfer( + self, + request: LlmRequest, + provider: AsyncTransferProvider = AsyncTransferProvider.GENERIC, + ) -> bool: """ Called after a send of KV cache is complete. 1. Decrements counter for request. @@ -478,7 +514,15 @@ def end_transfer(self, request: LlmRequest) -> bool: ) return False - if transfer_metadata.end_transfer(): + provider = AsyncTransferProvider(provider) + is_last_transfer = transfer_metadata.end_transfer(provider) + if is_last_transfer is None: + logger.warning( + f"Request {request.py_request_id} has no {provider.value} " + "transfer leg") + return False + + if is_last_transfer: self._requests_in_transfer.pop(request.py_request_id) self._request_transfer_metadata.pop(request.py_request_id) @@ -494,6 +538,19 @@ def end_transfer(self, request: LlmRequest) -> bool: return False + def has_transfer( + self, + request: LlmRequest, + provider: AsyncTransferProvider = AsyncTransferProvider.GENERIC, + ) -> bool: + """Return whether ``provider`` currently owns a request leg.""" + + transfer_metadata = self._request_transfer_metadata.get( + request.py_request_id) + if transfer_metadata is None: + return False + return transfer_metadata.has_transfer(AsyncTransferProvider(provider)) + def has_any_inflight_requests(self) -> bool: return len(self._requests_in_transfer) > 0 @@ -824,6 +881,8 @@ def __init__( self._error_budget = ErrorBudget() self._disagg_timed_out_ctx_cancelled_ids: set[int] = set() self._disagg_timed_out_gen_cancelled_ids: set[int] = set() + self._disagg_acknowledged_ctx_cancel_legs: set[int] = set() + self._disagg_peer_cancelled_ctx_ids: set[int] = set() self._disagg_inflight_cancel_unsupported_logged = False self.max_batch_size = max_batch_size self.adp_ctx_waiting_iters_count = 0 @@ -880,6 +939,10 @@ def on_detected(): # Waiting queue for requests that have been fetched but not yet scheduled self.waiting_queue: WaitingQueue = create_waiting_queue( waiting_queue_policy) + self._waiting_queue_policy = waiting_queue_policy + self._gen_first_pre_active_requests: Dict[int, LlmRequest] = {} + self._gen_first_pre_active_order: Dict[int, int] = {} + self._next_gen_first_pre_active_order = 0 self.control_request_barrier = threading.Event() self.control_action_done = threading.Event() @@ -1076,14 +1139,25 @@ def _maybe_init_kv_connector_manager(self): self.kv_connector_manager.wait_for_initialization() - def _end_transfer_and_maybe_terminate(self, request: LlmRequest): + def _end_transfer_and_maybe_terminate( + self, + request: LlmRequest, + provider: AsyncTransferProvider = AsyncTransferProvider.GENERIC, + ): transfer_failed = request.state == LlmRequestState.DISAGG_TRANS_ERROR + user_cancel_pending = (request.py_request_id in getattr( + self, "_disagg_acknowledged_ctx_cancel_legs", ()) + and self._request_vote_id(request) in getattr( + self, "canceled_req_ids", ())) if self.kv_cache_transceiver and request in self.active_requests: - if transfer_failed: + if transfer_failed or user_cancel_pending: # End only the transfer that just became terminal. Keep the - # request active so the synchronized error path can emit an - # error response after every async transfer releases ownership. - self.async_transfer_manager.end_transfer(request) + # request active so the synchronized error or user-cancel path + # can emit the terminal response after every async transfer + # releases ownership. In particular, a connector completion + # after the transceiver cancellation leg was acknowledged must + # not create a success response or terminate the request. + self.async_transfer_manager.end_transfer(request, provider) return # Fast-transfer: KV transfer completed in the same iteration # before _handle_responses could run. Create the response now @@ -1102,11 +1176,11 @@ def _end_transfer_and_maybe_terminate(self, request: LlmRequest): # participate. self._pending_transfer_responses.append( (request.py_request_id, response)) - if self.async_transfer_manager.end_transfer(request): + if self.async_transfer_manager.end_transfer(request, provider): self.active_requests.remove(request) self._terminate_request(request) return - if self.async_transfer_manager.end_transfer(request): + if self.async_transfer_manager.end_transfer(request, provider): if transfer_failed: return # Skip if the PP=1 early path already terminated this request; @@ -1468,6 +1542,17 @@ def shutdown(self): # no longer driving NCCL, so the send cannot deadlock. self._shutdown_sleep_wakeup_listeners() self.worker_started = False + shutdown_errors: List[Tuple[str, Exception]] = [] + + def run_shutdown_step(name: str, callback: Callable[[], None]) -> None: + try: + callback() + except Exception as error: + logger.error( + f"Shutdown step {name!r} failed: {error}\n{traceback.format_exc()}" + ) + shutdown_errors.append((name, error)) + # Release CUDA graphs before resource managers free their GPU memory. # Resource managers (e.g. SuffixAutomatonManager) allocate GPU workspace # that is referenced by raw pointers inside captured CUDA graphs. If @@ -1477,14 +1562,70 @@ def shutdown(self): # for the now-freed memory regions. for engine in (self.model_engine, self.draft_model_engine): if engine is not None and hasattr(engine, '_release_cuda_graphs'): - engine._release_cuda_graphs() + run_shutdown_step("release CUDA graphs", + engine._release_cuda_graphs) # Ensure graph destruction has fully completed on device before # resource managers start freeing GPU-backed workspaces. if torch.cuda.is_available(): - torch.cuda.synchronize() - for manager in self.resource_manager.resource_managers.values(): - if manager: - manager.shutdown() + run_shutdown_step("synchronize CUDA graph teardown", + torch.cuda.synchronize) + # Drain transfer workers and any dedicated consensus communicator + # before KV-cache managers release registered memory. The transceiver + # shutdown hook is idempotent and a no-op for implementations that do + # not own background resources. + transceiver_shutdown_terminal = self.kv_cache_transceiver is None + if self.kv_cache_transceiver is not None: + + def shutdown_cache_transceiver() -> None: + nonlocal transceiver_shutdown_terminal + primary_error: Optional[Exception] = None + max_attempts = 3 + deadline = time.monotonic() + 90.0 + for attempt in range(1, max_attempts + 1): + try: + result = self.kv_cache_transceiver.shutdown() + except Exception as error: + if primary_error is None: + primary_error = error + if attempt == max_attempts or time.monotonic( + ) >= deadline: + raise primary_error + continue + + if isinstance(result, threading.Event): + # An internal transfer progress thread cannot join + # itself. Registered memory remains owned until the + # event fires and a follow-up call confirms terminal + # teardown or surfaces its stored error. + remaining = deadline - time.monotonic() + if remaining <= 0 or not result.wait(remaining): + if primary_error is None: + primary_error = TimeoutError( + "timed out waiting for KV cache transceiver shutdown" + ) + raise primary_error + continue + + transceiver_shutdown_terminal = True + if primary_error is not None: + raise primary_error + return + + raise RuntimeError( + "KV cache transceiver shutdown did not reach a terminal state" + ) + + run_shutdown_step("KV cache transceiver", + shutdown_cache_transceiver) + if transceiver_shutdown_terminal: + for manager in self.resource_manager.resource_managers.values(): + if manager: + run_shutdown_step( + f"resource manager {type(manager).__name__}", + manager.shutdown) + else: + logger.error("Skipping resource-manager shutdown because KV cache " + "transceiver ownership did not reach a terminal state") # Note: do NOT call engine.cleanup() here. PyExecutor.shutdown() is # also invoked mid-init by configure_kv_cache_capacity() in # tensorrt_llm/_torch/pyexecutor/_util.py — the warmup pass calls @@ -1505,10 +1646,22 @@ def shutdown(self): # Stop the sampler's async worker, if it was used if (isinstance(self.sampler, AsyncWorkerMixin) and self.sampler.async_worker_enabled()): - self.sampler.async_worker_stop() + run_shutdown_step("sampler async worker", + self.sampler.async_worker_stop) if self.dwdp_manager is not None: - self.dwdp_manager.__exit__(None, None, None) + run_shutdown_step( + "DWDP manager", + lambda: self.dwdp_manager.__exit__(None, None, None), + ) self.dwdp_manager = None + if shutdown_errors: + primary_name, primary_error = shutdown_errors[0] + if len(shutdown_errors) > 1: + secondary_names = [name for name, _error in shutdown_errors[1:]] + logger.error( + f"Shutdown primary failure was {primary_name!r}; " + f"secondary failures occurred in {secondary_names}") + raise primary_error def can_enqueue_requests(self) -> bool: """ @@ -1601,7 +1754,8 @@ def set_gather_responses(self, gather_all_responses): @property def should_stop_processing(self): return self.is_shutdown and len(self.active_requests) == 0 and \ - len(self.waiting_queue) == 0 + len(self.waiting_queue) == 0 and \ + len(getattr(self, "_gen_first_pre_active_requests", {})) == 0 @contextmanager def _profiler(self): @@ -2431,6 +2585,8 @@ def _pp_schedule_and_propagate(self, microbatch_id: int): is_dp_broadcast = self.dist.tp_size > 1 and self.enable_attention_dp if self.dist.rank == 0 or (self.dist.is_first_pp_rank and is_dp_broadcast): + activated_context_request_ids = ( + self._activate_ready_pre_active_context_requests()) scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule( ) if self.kv_cache_transceiver: @@ -2438,8 +2594,11 @@ def _pp_schedule_and_propagate(self, microbatch_id: int): self._apply_disagg_transfer_admission( fitting_disagg_gen_init_requests)) serializable_schedule = SerializableSchedulerOutput.from_scheduler_result( - scheduled_batch, fitting_disagg_gen_init_requests, - num_fitting_reqs, wait_for_disagg_gen_transfer_progress) + scheduled_batch, + fitting_disagg_gen_init_requests, + num_fitting_reqs, + wait_for_disagg_gen_transfer_progress, + activated_context_request_ids=activated_context_request_ids) # Broadcast within first tp+cp group before send/recv chain to other tp+cp groups if self.dist.is_first_pp_rank: @@ -2468,11 +2627,23 @@ def _pp_schedule_and_propagate(self, microbatch_id: int): serializable_schedule, self.dist.next_pp_rank, PPCommTag.SCHEDULE_RESULT) + if scheduled_batch is None: + self._activate_pre_active_context_requests( + serializable_schedule.activated_context_request_ids) + if scheduled_batch is None: scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs = serializable_schedule.to_scheduler_result( self.active_requests) wait_for_disagg_gen_transfer_progress = ( serializable_schedule.wait_for_disagg_gen_transfer_progress) + if self.kv_cache_transceiver: + # The exact rank-zero PP schedule is the activation token for the + # asynchronous Python readiness protocol. Followers remain hidden + # until this point, immediately before their mirrored scheduler + # pass, so no rank can speculatively mutate scheduler/KV state for + # a request that rank zero did not select. + self.kv_cache_transceiver.activate_context_requests_for_schedule( + scheduled_batch.all_requests()) return (scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs, wait_for_disagg_gen_transfer_progress) @@ -2545,6 +2716,7 @@ def _executor_loop_pp(self): if self.kv_cache_transceiver: self._check_disagg_ctx_schedulable_status(new_requests) self._check_disagg_gen_transfer_status() + self._handle_pre_active_canceled_requests() if self.enable_iter_perf_stats: iter_stats = self._get_init_iter_stats( @@ -3739,7 +3911,8 @@ def _kv_connector_terminate_requests(self): if self.kv_connector_manager: reqs_to_terminate = self.kv_connector_manager.get_finished() for req in reqs_to_terminate: - self._end_transfer_and_maybe_terminate(req) + self._end_transfer_and_maybe_terminate( + req, AsyncTransferProvider.CONNECTOR) def _kv_connector_wait_for_save(self): if self.kv_connector_manager is not None: @@ -4265,15 +4438,20 @@ def _handle_control_request(self): pending = self.control_requests[0] - if pending.control_requires_drain and (len(self.active_requests) != 0 - or len(self.waiting_queue) != 0): + if pending.control_requires_drain and ( + len(self.active_requests) != 0 + or len(self.waiting_queue) != 0 or len( + getattr(self, "_gen_first_pre_active_requests", {})) != 0): # drain=True: keep the sentinel parked until the engine drains. return - logger.debug(f"[control_action] firing control request " - f"drain={pending.control_requires_drain} " - f"active_requests={len(self.active_requests)} " - f"waiting_queue={len(self.waiting_queue)}") + logger.debug( + f"[control_action] firing control request " + f"drain={pending.control_requires_drain} " + f"active_requests={len(self.active_requests)} " + f"waiting_queue={len(self.waiting_queue)} " + "pre_active_context_requests=" + f"{len(getattr(self, '_gen_first_pre_active_requests', {}))}") # Quiesce the device before the action mutates GPU state. Under the # overlap scheduler a previous batch's forward/sample kernels may still # be in flight, so an in-place update_weights reload (or sleep/wakeup @@ -4919,7 +5097,8 @@ def _fetch_and_enqueue_requests(self, waiting_queue: WaitingQueue, return # Calculate timeout - idle = (total_num_active_requests == 0) and len(waiting_queue) == 0 + idle = ((total_num_active_requests == 0) and len(waiting_queue) == 0 + and not getattr(self, "_gen_first_pre_active_requests", {})) if idle: # In Ray path (TLLM_DISABLE_MPI=1), use a periodic heartbeat timeout so rank 0 # reaches the broadcast path regularly to prevent trtllm-serve timeout when idle. @@ -4990,8 +5169,8 @@ def _pop_from_waiting_queue( @nvtx_range("_fetch_new_requests") def _fetch_new_requests( - self, waiting_queue: WaitingQueue, - active_requests: List[LlmRequest]) -> List[LlmRequest]: + self, waiting_queue: WaitingQueue, active_requests: List[LlmRequest] + ) -> Tuple[List[LlmRequest], List[LlmRequest]]: """Fetch new requests and return LlmRequests ready for execution.""" # 1. Gather rank states and calculate total_num_active_requests if self.enable_attention_dp: @@ -5035,6 +5214,8 @@ def _fetch_new_requests( self._fetch_and_enqueue_requests(waiting_queue, total_num_active_requests) + pre_active_items = self._take_pre_active_context_items(waiting_queue) + # 3. Pop requests from waiting queue new_requests = self._pop_from_waiting_queue( waiting_queue, total_num_active_requests, @@ -5042,10 +5223,11 @@ def _fetch_new_requests( # 4. Update performance metrics (before DP scheduling to clear all start_times) if self.enable_iter_perf_stats and self.dist.rank == 0: - self._update_new_active_requests_queue_latency(new_requests) + self._update_new_active_requests_queue_latency(new_requests + + pre_active_items) # 5. Update total fetch counter (used by benchmark disagg gating) - self.num_fetch_requests += len(new_requests) + self.num_fetch_requests += len(new_requests) + len(pre_active_items) # 6. Schedule requests across ranks (DP only) if self.enable_attention_dp: @@ -5070,12 +5252,110 @@ def _fetch_new_requests( new_requests = new_requests_cur_rank # 7. Merge requests - return merge_requests(new_requests, - cp_config=self.dist.cp_config, - cp_rank=self.dist.cp_rank, - cp_size=self.dist.cp_size, - exclude_last_generation_logits=self. - _should_exclude_last_generation_logits()) + merge_kwargs = dict(cp_config=self.dist.cp_config, + cp_rank=self.dist.cp_rank, + cp_size=self.dist.cp_size, + exclude_last_generation_logits=self. + _should_exclude_last_generation_logits()) + return (merge_requests(new_requests, **merge_kwargs), + merge_requests(pre_active_items, **merge_kwargs)) + + def _supports_pre_active_context_requests(self) -> bool: + return ( + self.kv_cache_transceiver is not None and + self.kv_cache_transceiver.supports_pre_active_context_requests() + is True) + + @staticmethod + def _is_pre_active_context_item(item: RequestQueueItem) -> bool: + request = item.request + params = getattr(request, "py_disaggregated_params", None) + # One queue item can materialize several child requests that do not + # have independent peer-metadata identities. Keep that uncommon path + # on the legacy capacity-accounted fallback until child activation can + # be represented atomically by the PP schedule. + return ( + item.is_normal_request and request is not None + and not item.child_req_ids + and request.request_type == RequestType.REQUEST_TYPE_CONTEXT_ONLY + and params is not None + and params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST) + + def _take_pre_active_context_items( + self, waiting_queue: WaitingQueue) -> List[RequestQueueItem]: + """Move qualified metadata waiters out of compute-active accounting.""" + if not self._supports_pre_active_context_requests(): + return [] + items = [ + item for item in waiting_queue + if self._is_pre_active_context_item(item) + ] + waiting_queue.remove_by_ids({item.id for item in items}) + return items + + def _pre_active_context_order_key(self, request: LlmRequest) -> Tuple: + request_id = request.request_id + insertion_order = self._gen_first_pre_active_order[request_id] + if self._waiting_queue_policy == WaitingQueuePolicy.PRIORITY: + return (-request.priority, insertion_order) + return (insertion_order, ) + + def _select_ready_pre_active_context_request_ids(self) -> List[int]: + if not self._supports_pre_active_context_requests(): + return [] + free_capacity = self.max_num_active_requests - len(self.active_requests) + if free_capacity <= 0: + return [] + ordered_requests = sorted(self._gen_first_pre_active_requests.values(), + key=self._pre_active_context_order_key) + return [ + request.request_id for request in ordered_requests + if self.kv_cache_transceiver. + is_context_request_ready_for_activation(request) + ][:free_capacity] + + def _activate_pre_active_context_requests(self, + request_ids: List[int]) -> None: + if not request_ids: + return + if len(set(request_ids)) != len(request_ids): + raise RuntimeError("duplicate pre-active context activation ID") + if len(self.active_requests) + len( + request_ids) > self.max_num_active_requests: + raise RuntimeError("pre-active context activation exceeds capacity") + missing_ids = [ + request_id for request_id in request_ids + if request_id not in self._gen_first_pre_active_requests + ] + if missing_ids: + raise RuntimeError( + "PP activation selected unknown pre-active context requests: " + f"{missing_ids}") + requests = [ + self._gen_first_pre_active_requests[request_id] + for request_id in request_ids + ] + unprepared_ids = [ + request.request_id for request in requests + if not self.kv_cache_transceiver. + is_context_request_ready_for_activation(request) + ] + if unprepared_ids: + raise RuntimeError( + "PP activation selected context requests without matching " + f"readiness state: {unprepared_ids}") + self.kv_cache_transceiver.activate_context_requests_for_schedule( + requests) + for request in requests: + request_id = request.request_id + del self._gen_first_pre_active_requests[request_id] + del self._gen_first_pre_active_order[request_id] + self.active_requests.append(request) + + def _activate_ready_pre_active_context_requests(self) -> List[int]: + request_ids = self._select_ready_pre_active_context_request_ids() + self._activate_pre_active_context_requests(request_ids) + return request_ids def _handle_special_queue_items( self, @@ -5114,7 +5394,9 @@ def _should_exclude_last_generation_logits(self) -> bool: def _fetch_and_activate_new_requests(self) -> List[LlmRequest]: - def _respond_if_invalid(request: LlmRequest) -> bool: + def _respond_if_invalid(request: LlmRequest, + *, + pre_active: bool = False) -> bool: """Immediately fail invalid request. Return True if invalid request was encountered and @@ -5124,21 +5406,61 @@ def _respond_if_invalid(request: LlmRequest) -> bool: self._validate_request(request) return False except Exception as e: - self._handle_errors(str(e), - requests=[request], - charge_budget=False) + if pre_active: + self._handle_pre_active_request_error(request, str(e)) + else: + self._handle_errors(str(e), + requests=[request], + charge_budget=False) return True - new_requests_cur_rank = self._fetch_new_requests( + new_requests_cur_rank, new_pre_active_requests = self._fetch_new_requests( self.waiting_queue, self.active_requests) validated_requests = [ request for request in new_requests_cur_rank if not _respond_if_invalid(request) ] + validated_pre_active_requests = [ + request for request in new_pre_active_requests + if not _respond_if_invalid(request, pre_active=True) + ] self.active_requests.extend(validated_requests) - return validated_requests + for request in validated_pre_active_requests: + request_id = request.request_id + if request_id in self._gen_first_pre_active_requests: + raise RuntimeError( + "duplicate generation-first pre-active request: " + f"{request_id}") + self._gen_first_pre_active_requests[request_id] = request + self._gen_first_pre_active_order[ + request_id] = self._next_gen_first_pre_active_order + self._next_gen_first_pre_active_order += 1 + canceled_request_ids = set(self.canceled_req_ids) + canceled_pre_active_requests = [ + request for request in validated_pre_active_requests + if self._request_vote_id(request) in canceled_request_ids + ] + if canceled_pre_active_requests: + # Request and cancellation queue items can arrive in one fetch. + # Keep those requests on the no-resource lane, but establish the + # cancellation marker before readiness can publish a vote. + self.kv_cache_transceiver.exclude_context_requests_from_readiness( + canceled_pre_active_requests) + pending_pre_active_requests = [ + request for request in validated_pre_active_requests + if self._request_vote_id(request) not in canceled_request_ids + ] + if pending_pre_active_requests: + self.kv_cache_transceiver.prepare_context_requests( + pending_pre_active_requests) + # Preserve the established per-iteration admission accounting: before + # this lane existed, generation-first metadata waiters entered + # active_requests and were counted as newly admitted here. They no + # longer consume compute capacity, but they have left the request queue + # and their queue-latency start time was consumed above. + return validated_requests + validated_pre_active_requests def _add_kv_cache_events(self): kv_cache_manager = self.resource_manager.resource_managers.get( @@ -5621,7 +5943,9 @@ def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None: req.py_kv_transfer_timed_out = True for req in self.async_transfer_manager.requests_in_transfer().values(): - flag_if_kv_transfer_timed_out(req, "context") + if self.async_transfer_manager.has_transfer( + req, AsyncTransferProvider.TRANSCEIVER): + flag_if_kv_transfer_timed_out(req, "context") for req in self.active_requests: if req.is_disagg_generation_transmission_in_progress: @@ -5638,10 +5962,30 @@ def _check_disagg_ctx_schedulable_status(self, """ if not self.kv_cache_transceiver: return + canceled_request_ids = set(self.canceled_req_ids) + canceled_gen_first_ctx_requests = [ + req for req in ( + *self.active_requests, + *getattr(self, "_gen_first_pre_active_requests", {}).values()) + if req.is_context_only_request and req.py_disaggregated_params. + schedule_style == DisaggScheduleStyle.GENERATION_FIRST + and self._request_vote_id(req) in canceled_request_ids + ] + exclude_from_readiness = getattr( + self.kv_cache_transceiver, + "exclude_context_requests_from_readiness", + None, + ) + if canceled_gen_first_ctx_requests and callable(exclude_from_readiness): + # Cancellation queue items are known before readiness progresses. + # Publish withdrawals first so prepare/ACK/release cannot make a + # cancelled request schedulable in this iteration. + exclude_from_readiness(canceled_gen_first_ctx_requests) gen_first_ctx_requests = [ req for req in new_requests if req.is_context_only_request and req.py_disaggregated_params. schedule_style == DisaggScheduleStyle.GENERATION_FIRST + and self._request_vote_id(req) not in canceled_request_ids ] # Always call prepare_context_requests when there are new requests # or previously-waiting requests, so the tp_allgather consensus @@ -6064,7 +6408,8 @@ def kv_connector_request_finished(req: LlmRequest): else: if self.kv_connector_manager.request_finished( req, cache_block_ids): - self.async_transfer_manager.start_transfer(req) + self.async_transfer_manager.start_transfer( + req, AsyncTransferProvider.CONNECTOR) if self.kv_cache_transceiver: for req in scheduled_requests: @@ -6079,7 +6424,8 @@ def kv_connector_request_finished(req: LlmRequest): req.py_request_id) # Order is important here: we need to start the transfer before responding # to make sure the blocks are stored for reuse before they are sent. - self.async_transfer_manager.start_transfer(req) + self.async_transfer_manager.start_transfer( + req, AsyncTransferProvider.TRANSCEIVER) self.kv_cache_transceiver.respond_and_send_async(req) if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: @@ -6142,7 +6488,26 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): finished_requests, error_requests = self.kv_cache_transceiver.check_context_transfer_status( atLeastNum) + take_cancelled_ids = getattr( + self.kv_cache_transceiver, + "take_context_cancelled_request_ids", + None, + ) + newly_cancelled_ids = [] + if callable(take_cancelled_ids): + peer_cancelled_ids = getattr(self, "_disagg_peer_cancelled_ctx_ids", + None) + if peer_cancelled_ids is None: + peer_cancelled_ids = set() + self._disagg_peer_cancelled_ctx_ids = peer_cancelled_ids + newly_cancelled_ids = list(take_cancelled_ids()) + peer_cancelled_ids.update(newly_cancelled_ids) + completed_req_ids = set(finished_requests + error_requests) + # A newly surfaced cancellation is terminal for timeout suppression, + # but it must not pass through the ordinary completion loop: its state + # and finish reason are assigned by the peer-cancellation path below. + terminal_req_ids = completed_req_ids | set(newly_cancelled_ids) requests_in_transfer = self.async_transfer_manager.requests_in_transfer( ) @@ -6156,7 +6521,59 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): request = requests_in_transfer[request_id] - self._end_transfer_and_maybe_terminate(request) + if not self.async_transfer_manager.has_transfer( + request, AsyncTransferProvider.TRANSCEIVER): + logger.warning( + f"Ignoring duplicate transceiver completion for request " + f"{request_id}; only a non-transceiver transfer leg remains" + ) + continue + self._end_transfer_and_maybe_terminate( + request, AsyncTransferProvider.TRANSCEIVER) + + # A peer/global cancellation is terminal only after the transceiver + # reports local quiescence. Acknowledge its request-scoped tombstone, + # retire exactly the transceiver's AsyncTransferManager leg, and route + # it through the ordinary transfer-error path. A locally queued user + # cancellation keeps its CANCELLED finish reason and is completed by + # _handle_canceled_requests instead. + local_user_cancelled_ids = set(self.canceled_req_ids) + peer_cancelled_ids = getattr(self, "_disagg_peer_cancelled_ctx_ids", + set()) + for request_id in list(peer_cancelled_ids): + requests_in_transfer = self.async_transfer_manager.requests_in_transfer( + ) + request = requests_in_transfer.get(request_id) + if request is None: + request = next( + (active_request for active_request in self.active_requests + if active_request.py_request_id == request_id), + None, + ) + if request is None: + logger.warning( + f"Peer-cancelled request {request_id} was not found in " + "the transfer manager or active request set; retaining " + "the cancellation handoff for retry") + continue + if self._request_vote_id(request) in local_user_cancelled_ids: + peer_cancelled_ids.discard(request_id) + continue + if not self.async_transfer_manager.has_transfer( + request, AsyncTransferProvider.TRANSCEIVER): + # Preserve the globally committed terminal decision without + # consuming the connector's independently owned leg. Error + # cleanup remains deferred until the connector callback. + request.state = LlmRequestState.DISAGG_TRANS_ERROR + peer_cancelled_ids.discard(request_id) + continue + if not self._request_kv_transfer_cancellation(request): + continue + request.state = LlmRequestState.DISAGG_TRANS_ERROR + if request_id in requests_in_transfer: + self._end_transfer_and_maybe_terminate( + request, AsyncTransferProvider.TRANSCEIVER) + peer_cancelled_ids.discard(request_id) # The set of requests in transfer may have changed since we terminated some requests. requests_in_transfer = self.async_transfer_manager.requests_in_transfer( @@ -6164,8 +6581,11 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): for request_id in list(requests_in_transfer.keys()): request = requests_in_transfer[request_id] + if not self.async_transfer_manager.has_transfer( + request, AsyncTransferProvider.TRANSCEIVER): + continue if (not request.py_kv_transfer_timed_out - or request_id in completed_req_ids + or request_id in terminal_req_ids or request_id in self._disagg_timed_out_ctx_cancelled_ids): continue @@ -6184,7 +6604,8 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): # cancelled is immediately released from the async manager. request.py_kv_transfer_start_time = None request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE - self._end_transfer_and_maybe_terminate(request) + self._end_transfer_and_maybe_terminate( + request, AsyncTransferProvider.TRANSCEIVER) self._check_cache_transfer_errors("context requests") @@ -6545,8 +6966,12 @@ def _handle_errors(self, f"Drained {len(waiting_responses)} queued requests " "on fatal error") - failed_requests = (list(self.active_requests) - if requests is None else requests) + pre_active_request_ids = set( + getattr(self, "_gen_first_pre_active_requests", {})) + failed_requests = ( + (list(self.active_requests) + + list(getattr(self, "_gen_first_pre_active_requests", {}).values())) + if requests is None else requests) for request in failed_requests: req_id = request.py_request_id request.state = LlmRequestState.GENERATION_COMPLETE @@ -6556,14 +6981,25 @@ def _handle_errors(self, client_id=request.py_client_id) if requests is None: self.active_requests.clear() + getattr(self, "_gen_first_pre_active_requests", {}).clear() + getattr(self, "_gen_first_pre_active_order", {}).clear() else: self.active_requests = [ request for request in self.active_requests if request not in requests ] + for request in requests: + request_id = request.request_id + getattr(self, "_gen_first_pre_active_requests", + {}).pop(request_id, None) + getattr(self, "_gen_first_pre_active_order", + {}).pop(request_id, None) self._enqueue_responses(list(error_responses.items())) for request in failed_requests: - self._terminate_request(request) + if request.request_id in pre_active_request_ids: + self._release_pre_active_request_bookkeeping(request) + else: + self._terminate_request(request) if self._fatal_error is not None: self.executor_request_queue.enqueue_shutdown_request() @@ -6584,10 +7020,41 @@ def _do_terminate_request(self, request: LlmRequest): self._prefetched_request_ids.discard(request.py_request_id) self._disagg_timed_out_ctx_cancelled_ids.discard(request.py_request_id) self._disagg_timed_out_gen_cancelled_ids.discard(request.py_request_id) + getattr(self, "_disagg_acknowledged_ctx_cancel_legs", + set()).discard(request.py_request_id) + getattr(self, "_disagg_peer_cancelled_ctx_ids", + set()).discard(request.py_request_id) + + if self.gather_all_responses or self.dist.rank == 0: + self.result_wait_queues.pop(request.py_request_id, None) + def _release_pre_active_request_bookkeeping(self, + request: LlmRequest) -> None: + """Release bookkeeping for a request that never owned model resources.""" + getattr(self, "_gen_first_pre_active_requests", + {}).pop(request.request_id, None) + getattr(self, "_gen_first_pre_active_order", + {}).pop(request.request_id, None) + self._prefetched_request_ids.discard(request.py_request_id) + self._disagg_timed_out_ctx_cancelled_ids.discard(request.py_request_id) + self._disagg_timed_out_gen_cancelled_ids.discard(request.py_request_id) + getattr(self, "_disagg_acknowledged_ctx_cancel_legs", + set()).discard(request.py_request_id) + getattr(self, "_disagg_peer_cancelled_ctx_ids", + set()).discard(request.py_request_id) if self.gather_all_responses or self.dist.rank == 0: self.result_wait_queues.pop(request.py_request_id, None) + def _handle_pre_active_request_error(self, request: LlmRequest, + error_msg: str) -> None: + """Fail a request that was validated before resource activation.""" + request.state = LlmRequestState.GENERATION_COMPLETE + response = LlmResponse(request_id=request.py_request_id, + error_msg=error_msg, + client_id=request.py_client_id) + self._enqueue_responses([(request.py_request_id, response)]) + self._release_pre_active_request_bookkeeping(request) + def _is_request_in_transmission(self, request) -> bool: """Check if a request is currently in transmission state.""" return (request.state @@ -6605,14 +7072,60 @@ def _try_cancel_request(self, request) -> bool: return True async_transfer_manager = getattr(self, "async_transfer_manager", None) - if (getattr(request, "is_context_only_request", False) is True - and async_transfer_manager is not None and request.py_request_id - in async_transfer_manager.requests_in_transfer()): - if self._is_disagg_inflight_cancel_active(): - self._request_kv_transfer_cancellation(request) + is_context_request = (getattr(request, "is_context_only_request", False) + is True) + requests_in_transfer = (async_transfer_manager.requests_in_transfer() + if is_context_request + and async_transfer_manager is not None else {}) + acknowledged_legs = getattr(self, + "_disagg_acknowledged_ctx_cancel_legs", + None) + if is_context_request and async_transfer_manager is not None: + if acknowledged_legs is None: + acknowledged_legs = set() + self._disagg_acknowledged_ctx_cancel_legs = acknowledged_legs + request_id = request.py_request_id + if request_id in acknowledged_legs: + if request_id not in requests_in_transfer: + # The transceiver leg was acknowledged on an earlier + # cancellation pass and the final connector leg is now + # quiescent. Do not retry transceiver cancellation based + # on a stale TRANS_IN_PROGRESS request state. + acknowledged_legs.discard(request_id) + return True + # Another asynchronous transfer (for example, a cache + # connector save) still owns the request. Do not decrement its + # transfer-manager leg while retrying the user cancellation. + return False + + transceiver_owns_request = self.kv_cache_transceiver.owns_request( + request) is True + + if (is_context_request and async_transfer_manager is not None + and request.py_request_id in requests_in_transfer): + request_id = request.py_request_id + transceiver_leg_inflight = async_transfer_manager.has_transfer( + request, AsyncTransferProvider.TRANSCEIVER) is True + if not transceiver_leg_inflight: + # A connector can outlive a normally completed transceiver. + # A process-global cancellation capability (or a no-session + # transceiver cancel that reports success) is not evidence + # that the remaining connector leg is quiescent. + return False + if (transceiver_owns_request + or self._is_disagg_inflight_cancel_active() + ) and self._request_kv_transfer_cancellation(request): + assert acknowledged_legs is not None + acknowledged_legs.add(request_id) + if async_transfer_manager.end_transfer( + request, AsyncTransferProvider.TRANSCEIVER): + acknowledged_legs.discard(request_id) + return True return False if not self._is_request_in_transmission(request): + if transceiver_owns_request: + return self._request_kv_transfer_cancellation(request) return True if self._is_disagg_inflight_cancel_active(): @@ -6632,7 +7145,19 @@ def _handle_canceled_requests(self): # Remove canceled requests from the waiting queue self.waiting_queue.remove_by_ids(canceled_req_ids_set) - still_pending_canceled_ids = [] + pre_active_vote_ids = { + request.py_request_id + if not request.is_child else request.parent_request_id + for request in getattr(self, "_gen_first_pre_active_requests", + {}).values() + } + still_pending_canceled_ids = [ + request_id for request_id in self.canceled_req_ids + if request_id in pre_active_vote_ids + ] + # Metadata-only pre-active requests are drained by + # _handle_pre_active_canceled_requests without touching model resource + # managers. Never route them through this active-request cleanup path. for request in self.active_requests: req_id = request.py_request_id if not request.is_child else request.parent_request_id if req_id not in canceled_req_ids_set: @@ -6646,12 +7171,68 @@ def _handle_canceled_requests(self): request.finish_by_reason(FinishReason.CANCELLED) request.decoding_iter = request.py_decoding_iter else: - still_pending_canceled_ids.append(req_id) + if req_id not in still_pending_canceled_ids: + still_pending_canceled_ids.append(req_id) # Clear list of requests marked for cancellation and add back those that failed to cancel. self.canceled_req_ids.clear() self.canceled_req_ids.extend(still_pending_canceled_ids) + def _handle_pre_active_canceled_requests(self) -> None: + """Drain metadata-only CTX cancellation without an executed batch. + + PP normally processes cancellation while retiring an executed batch. + A generation-first metadata waiter intentionally owns no executable + batch or model resource, so that callback may never run. Retire only + pre-active requests here; ordinary active cancellation keeps its + existing batch-synchronized lifecycle. + """ + pre_active_requests = getattr(self, "_gen_first_pre_active_requests", + {}) + if not pre_active_requests: + return + if self.is_shutdown: + queued_cancellations = set(self.canceled_req_ids) + for request in pre_active_requests.values(): + request_id = (request.py_request_id if not request.is_child else + request.parent_request_id) + if request_id not in queued_cancellations: + self.canceled_req_ids.append(request_id) + queued_cancellations.add(request_id) + if not self.canceled_req_ids: + return + canceled_ids = set(self.canceled_req_ids) + completed_ids = set() + responses = [] + for request in list(pre_active_requests.values()): + request_id = (request.py_request_id if not request.is_child else + request.parent_request_id) + if request_id not in canceled_ids: + continue + if not self._try_cancel_request(request): + continue + + request.py_kv_transfer_timed_out = False + request.finish_by_reason(FinishReason.CANCELLED) + request.decoding_iter = request.py_decoding_iter + request.draft_tokens = request.py_draft_tokens or [] + response = request.create_response(False, self.dist.rank) + if response is not None: + response.result.cached_tokens = request.cached_tokens + self._maybe_attach_ctx_usage(request, response) + responses.append((request.py_request_id, response)) + + self._release_pre_active_request_bookkeeping(request) + completed_ids.add(request_id) + + if responses: + self._enqueue_responses(responses) + if completed_ids: + self.canceled_req_ids[:] = [ + request_id for request_id in self.canceled_req_ids + if request_id not in completed_ids + ] + @nvtx_range("_enqueue_responses") def _enqueue_responses(self, responses: Iterable[Tuple[int, LlmResponse]]): if 0 not in self.dist.mapping.tp_group and not self.gather_all_responses: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 7920c8942e71..f96136444a4e 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -1016,6 +1016,11 @@ def drafting_loop_wrapper(model): peft_cache_config=peft_cache_config, scheduler_config=scheduler_config, cache_transceiver_config=cache_transceiver_config, + # The estimation executor must retain the transceiver's registered + # memory so KV-capacity accounting matches the final runtime. It + # is not, however, a serving lifetime: never let its temporary + # endpoints escape through get_disaggregated_params(). + publish_disaggregated_params=not estimating_kv_cache, virtual_memory_pools=vm_pools if not estimating_kv_cache else None, execution_stream=execution_stream, max_num_sequences=max_num_seq_slots, @@ -1095,6 +1100,7 @@ def drafting_loop_wrapper(model): execution_stream=execution_stream, dwdp_manager=dwdp_manager, max_num_sequences=max_num_seq_slots, + publish_disaggregated_params=True, ) _adjust_torch_mem_fraction() diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index caa8e3cb3de1..3e5678ba7a00 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import dataclasses import inspect from abc import ABC, abstractmethod @@ -257,6 +260,7 @@ class SerializableSchedulerOutput: ] # request ids of fitting disaggregated generation initialization requests num_fitting_requests: int # number of fitting requests wait_for_disagg_gen_transfer_progress: bool = False + activated_context_request_ids: list[int] = dataclasses.field(default_factory=list) @classmethod def from_scheduler_result( @@ -265,6 +269,7 @@ def from_scheduler_result( fitting_disagg_gen_init_requests: RequestList, num_fitting_requests: int, wait_for_disagg_gen_transfer_progress: bool = False, + activated_context_request_ids: Optional[list[int]] = None, ) -> "SerializableSchedulerOutput": return cls( encoder_requests=[req.request_id for req in scheduled_requests.encoder_requests], @@ -281,6 +286,7 @@ def from_scheduler_result( ], num_fitting_requests=num_fitting_requests, wait_for_disagg_gen_transfer_progress=wait_for_disagg_gen_transfer_progress, + activated_context_request_ids=activated_context_request_ids or [], ) def to_scheduler_result( diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 577717586748..ceee459fd4c7 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -163,7 +163,6 @@ def __init__(self, self._executor_cls = kwargs.pop("executor_cls", GenerationExecutor) self._orchestrator_type = kwargs.get("orchestrator_type", None) self._llm_id = None - self._disaggregated_params: Optional[dict] = None log_level = logger.level logger.set_level("info") # force display the backend @@ -342,10 +341,11 @@ def get_data_transceiver_state(self) -> bytes: @property @set_api_status("beta") def disaggregated_params(self) -> dict: - if self._disaggregated_params is None: - self._disaggregated_params = self._executor.get_disaggregated_params( - ) if self._executor else {} - return self._disaggregated_params + # Worker initialization has an internal KV-capacity estimation phase. + # Never freeze the first RPC result: only the current, final executor + # owns endpoints that may be advertised to a generation-first peer. + return self._executor.get_disaggregated_params( + ) if self._executor else {} @staticmethod def _is_token_id_list(value: Any) -> bool: diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index e2ee00257e98..80299eabe2de 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -266,6 +266,14 @@ def _get_gen_request( if ctx_server_info and "server_info" in ctx_server_info: disaggregated_params = ctx_server_info["server_info"].get("disaggregated_params", {}) if disaggregated_params: + # Used only by the coordinator to reject stale endpoint + # lifetimes; never forward this internal ownership token in an + # inference request. + disaggregated_params = { + key: value + for key, value in disaggregated_params.items() + if key != "ctx_endpoint_generation" + } # ctx_info_endpoint from get_disaggregated_params() is a list; # the Pydantic model expects a single str. ep = disaggregated_params.get("ctx_info_endpoint") @@ -399,6 +407,51 @@ async def _send_disagg_request_gen_first( ctx_server, ctx_server_info = await self._ctx_router.get_next_server( request, req_id=disagg_request_id ) + cached_disaggregated_params = ctx_server_info.get("server_info", {}).get( + "disaggregated_params", {} + ) + cached_endpoint = cached_disaggregated_params.get("ctx_info_endpoint") + cached_generation = cached_disaggregated_params.get("ctx_endpoint_generation") + if not cached_endpoint or cached_generation: + try: + runtime_server_info = await self._ctx_router.get_runtime_server_info( + ctx_server, + require_generation=bool(cached_generation), + ) + except BaseException: + # Selection has already reserved router/coordinator load, + # but no OpenAIClient owns it yet. Release it here and + # shield cleanup from request cancellation. + cleanup_task = asyncio.create_task( + self._ctx_router.finish_request( + request, + success=False, + req_id=disagg_request_id, + ) + ) + + def report_cleanup_error(completed_task: asyncio.Task) -> None: + if completed_task.cancelled(): + return + cleanup_error = completed_task.exception() + if cleanup_error is not None: + logger.error( + "Failed to release context routing " + "reservation after server-info refresh " + f"failure: {cleanup_error}" + ) + + cleanup_task.add_done_callback(report_cleanup_error) + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + pass + except Exception: + # The callback above records cleanup failure without + # replacing the original refresh exception. + pass + raise + ctx_server_info = {"server_info": runtime_server_info} ctx_req = self._get_ctx_request(request, disagg_request_id) gen_req = self._get_gen_request( request, diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index f0588a7e82f6..927692974600 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -44,6 +44,9 @@ COORDINATOR_FINISH_WORKERS = 16 COORDINATOR_FINISH_QUEUE_SIZE = 4096 COORDINATOR_FINISH_DRAIN_TIMEOUT_S = 5.0 +RUNTIME_SERVER_INFO_MAX_ATTEMPTS = 3 +RUNTIME_SERVER_INFO_RETRY_DELAY_S = 0.1 +RUNTIME_SERVER_INFO_REQUEST_TIMEOUT_S = 5.0 # Max number of conversations whose home-server pin is retained (LRU). ROUTE_AFFINITY_CACHE_SIZE = 50000 @@ -344,9 +347,19 @@ def __init__( self._health_check_timeout = metadata_server_cfg.health_check_timeout if metadata_server_cfg else None self._server_preparation_func = server_preparation_func self._prepared_ready_servers: set[str] = set() + self._runtime_server_info_refresh_tasks: dict[tuple[str, bool], + asyncio.Task] = {} + self._runtime_server_info_logged_generations: dict[str, str] = {} async def close(self): """Close the shared HTTP session.""" + refresh_tasks = list(self._runtime_server_info_refresh_tasks.values()) + self._runtime_server_info_refresh_tasks.clear() + self._runtime_server_info_logged_generations.clear() + for task in refresh_tasks: + task.cancel() + if refresh_tasks: + await asyncio.gather(*refresh_tasks, return_exceptions=True) if self._session: try: await self._session.close() @@ -416,6 +429,117 @@ async def _prepare_server(self, server: str): # swallow the error, if the server becomes ready or is added later, it will be prepared again logger.warning(f"Error preparing server {server}: {e}") + async def _fetch_runtime_server_info(self, server: str, + require_generation: bool) -> dict: + last_error: Optional[Exception] = None + for attempt in range(RUNTIME_SERVER_INFO_MAX_ATTEMPTS): + try: + server_info = await self._fetch_server_info( + server, self._health_check_timeout + or RUNTIME_SERVER_INFO_REQUEST_TIMEOUT_S) + disaggregated_params = server_info.get("disaggregated_params", + {}) + endpoint = disaggregated_params.get("ctx_info_endpoint") + generation = disaggregated_params.get("ctx_endpoint_generation") + if not endpoint or (require_generation and not generation): + raise RuntimeError("server has not published final-runtime " + "disaggregated endpoint ownership") + return server_info + except RuntimeError as error: + last_error = error + if attempt + 1 < RUNTIME_SERVER_INFO_MAX_ATTEMPTS: + await asyncio.sleep(RUNTIME_SERVER_INFO_RETRY_DELAY_S) + + raise RuntimeError( + f"Failed to fetch final-runtime server info for {server}" + ) from last_error + + async def get_runtime_server_info( + self, + server: str, + require_generation: bool = False, + validate_server_membership: bool = True) -> dict: + """Return endpoint metadata owned by the final serving executor. + + Startup preparation can overlap remote worker construction. Coalesce + concurrent generation-first refreshes per server and accept a result + only when the Python transceiver reports an endpoint and, when the + asynchronous-readiness path advertised one, its lifetime UUID. Every + later asynchronous dispatch fetches again, so a same-URL restart + cannot leave a retired endpoint cached. Network waits happen outside + the router lock so routing and monitoring remain responsive. + """ + refresh_key = (server, require_generation) + async with self._lock: + refresh_task = self._runtime_server_info_refresh_tasks.get( + refresh_key) + if refresh_task is None: + refresh_task = asyncio.create_task( + self._fetch_runtime_server_info(server, require_generation)) + self._runtime_server_info_refresh_tasks[ + refresh_key] = refresh_task + + def retire_refresh(completed_task: asyncio.Task) -> None: + # Event-loop callbacks are serialized. Retire the task + # independently of its waiters so cancellation cannot + # leave a completed, stale refresh cached. + if self._runtime_server_info_refresh_tasks.get( + refresh_key) is completed_task: + self._runtime_server_info_refresh_tasks.pop( + refresh_key, None) + if not completed_task.cancelled(): + completed_task.exception() + + refresh_task.add_done_callback(retire_refresh) + + try: + server_info = await asyncio.shield(refresh_task) + except BaseException: + if refresh_task.done(): + async with self._lock: + if self._runtime_server_info_refresh_tasks.get( + refresh_key) is refresh_task: + self._runtime_server_info_refresh_tasks.pop( + refresh_key, None) + raise + + async with self._lock: + if validate_server_membership and server not in self._servers: + raise RuntimeError( + f"Server {server} was removed while refreshing runtime info" + ) + if not validate_server_membership: + # A delegating client receives coordinator-selected workers + # without mirroring the coordinator's dynamic server list. + # Return their live metadata to this request, but do not retain + # churned fleet URLs in the wrapped local router indefinitely. + if self._runtime_server_info_refresh_tasks.get( + refresh_key) is refresh_task: + self._runtime_server_info_refresh_tasks.pop( + refresh_key, None) + return server_info + previous_generation = self._server_info.get(server, {}).get( + "disaggregated_params", {}).get("ctx_endpoint_generation") + current_generation = server_info["disaggregated_params"].get( + "ctx_endpoint_generation") + if (previous_generation is not None + and previous_generation != current_generation): + logger.info(f"server {server} endpoint generation changed from " + f"{previous_generation} to {current_generation}") + if (current_generation is not None + and self._runtime_server_info_logged_generations.get(server) + != current_generation): + logger.info("PYTHON_ASYNC_CONSENSUS " + "transition=runtime_endpoint_refresh " + f"server={server} generation={current_generation}") + self._runtime_server_info_logged_generations[ + server] = current_generation + self._server_info[server] = server_info + if self._runtime_server_info_refresh_tasks.get( + refresh_key) is refresh_task: + self._runtime_server_info_refresh_tasks.pop(refresh_key, None) + return server_info + async def prepare_servers(self, servers: Optional[List[str]] = None): targets = self._servers if servers is None else servers for server in targets: @@ -468,6 +592,13 @@ async def remove_server(self, server: str): ] self._on_servers_updated(old_servers, self._servers) self._prepared_ready_servers.discard(server) + self._runtime_server_info_logged_generations.pop(server, None) + refresh_keys = [ + key for key in self._runtime_server_info_refresh_tasks + if key[0] == server + ] + for key in refresh_keys: + self._runtime_server_info_refresh_tasks.pop(key).cancel() self._server_info.pop(server, None) logger.debug( f"Removed server {server}, current server list: {self._servers}") @@ -548,6 +679,16 @@ async def _monitor_servers(self, poll_interval: float = 10.0): for server in old_servers: if server not in final_servers: self._prepared_ready_servers.discard(server) + self._runtime_server_info_logged_generations.pop( + server, None) + refresh_keys = [ + key for key in + self._runtime_server_info_refresh_tasks + if key[0] == server + ] + for key in refresh_keys: + self._runtime_server_info_refresh_tasks.pop( + key).cancel() self._server_info.pop(server, None) logger.info(f"Server {server} is removed") @@ -1700,6 +1841,21 @@ def session(self) -> aiohttp.ClientSession: def _on_servers_updated(self, old_servers, new_servers): pass + async def get_runtime_server_info( + self, + server: str, + require_generation: bool = False, + validate_server_membership: bool = True) -> dict: + # The proxy session targets the coordinator (and may use a UDS). The + # selected worker's /server_info must be queried through the wrapped + # local router and its ordinary HTTP session. + del validate_server_membership + return await self._local.get_runtime_server_info( + server, + require_generation=require_generation, + validate_server_membership=False, + ) + def _request_id(self, request: OpenAIRequest, req_id: Optional[int] = None) -> int: diff --git a/tests/integration/defs/perf/disagg_server_config.py b/tests/integration/defs/perf/disagg_server_config.py new file mode 100644 index 000000000000..3aaf8fcd8e29 --- /dev/null +++ b/tests/integration/defs/perf/disagg_server_config.py @@ -0,0 +1,60 @@ +# 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. + +"""Helpers for generating perf-sanity disaggregated-server configs.""" + +from typing import Any, Mapping, Optional + +_ALLOWED_SCHEDULE_STYLES = frozenset({"context_first", "generation_first"}) + + +def build_disagg_server_config( + hostname: str, + port: int, + num_ctx_servers: int, + num_gen_servers: int, + ctx_hostnames: list[str], + gen_hostnames: list[str], + server_config_extra: Optional[Mapping[str, Any]] = None, +) -> dict[str, Any]: + """Build the runtime config with an optional schedule-style override.""" + extras = dict(server_config_extra or {}) + unsupported_keys = sorted(set(extras) - {"schedule_style"}) + if unsupported_keys: + raise ValueError( + "server_config_extra supports only 'schedule_style'; " + f"unsupported keys: {unsupported_keys}" + ) + if "schedule_style" in extras and extras["schedule_style"] not in _ALLOWED_SCHEDULE_STYLES: + raise ValueError( + "server_config_extra.schedule_style must be one of " + f"{sorted(_ALLOWED_SCHEDULE_STYLES)}, got {extras['schedule_style']!r}" + ) + + server_config = { + "hostname": hostname, + "port": port, + "backend": "pytorch", + "context_servers": { + "num_instances": num_ctx_servers, + "urls": ctx_hostnames, + }, + "generation_servers": { + "num_instances": num_gen_servers, + "urls": gen_hostnames, + }, + } + server_config.update(extras) + return server_config diff --git a/tests/integration/defs/perf/test_disagg_server_config.py b/tests/integration/defs/perf/test_disagg_server_config.py new file mode 100644 index 000000000000..ff68cdfe7092 --- /dev/null +++ b/tests/integration/defs/perf/test_disagg_server_config.py @@ -0,0 +1,107 @@ +# 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. + +"""Unit tests for perf-sanity disaggregated-server config generation.""" + +import unittest + +try: + from .disagg_server_config import build_disagg_server_config +except ImportError: + # Support direct execution without importing the GPU-heavy integration package. + from disagg_server_config import build_disagg_server_config + + +class TestBuildDisaggServerConfig(unittest.TestCase): + def test_no_extra_preserves_baseline_config(self): + config = build_disagg_server_config( + hostname="benchmark-host", + port=8123, + num_ctx_servers=1, + num_gen_servers=2, + ctx_hostnames=["ctx-host:8000"], + gen_hostnames=["gen-host-0:8001", "gen-host-1:8001"], + ) + + self.assertEqual( + config, + { + "hostname": "benchmark-host", + "port": 8123, + "backend": "pytorch", + "context_servers": { + "num_instances": 1, + "urls": ["ctx-host:8000"], + }, + "generation_servers": { + "num_instances": 2, + "urls": ["gen-host-0:8001", "gen-host-1:8001"], + }, + }, + ) + + def test_applies_schedule_style_from_server_config_extra(self): + config = build_disagg_server_config( + hostname="benchmark-host", + port=8123, + num_ctx_servers=1, + num_gen_servers=1, + ctx_hostnames=["ctx-host:8000"], + gen_hostnames=["gen-host:8001"], + server_config_extra={"schedule_style": "generation_first"}, + ) + + self.assertEqual(config["schedule_style"], "generation_first") + self.assertEqual(config["context_servers"]["urls"], ["ctx-host:8000"]) + self.assertEqual(config["generation_servers"]["urls"], ["gen-host:8001"]) + + def test_rejects_unknown_or_reserved_extra_keys(self): + for extra in ( + {"unexpected": "value"}, + {"hostname": "redirected-host"}, + {"port": 9000}, + {"backend": "other"}, + {"context_servers": {}}, + {"generation_servers": {}}, + ): + with ( + self.subTest(extra=extra), + self.assertRaisesRegex(ValueError, "supports only 'schedule_style'"), + ): + build_disagg_server_config( + hostname="benchmark-host", + port=8123, + num_ctx_servers=1, + num_gen_servers=1, + ctx_hostnames=["ctx-host:8000"], + gen_hostnames=["gen-host:8001"], + server_config_extra=extra, + ) + + def test_rejects_invalid_schedule_style(self): + with self.assertRaisesRegex(ValueError, "schedule_style must be one of"): + build_disagg_server_config( + hostname="benchmark-host", + port=8123, + num_ctx_servers=1, + num_gen_servers=1, + ctx_hostnames=["ctx-host:8000"], + gen_hostnames=["gen-host:8001"], + server_config_extra={"schedule_style": "unsupported"}, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/integration/defs/perf/test_perf_sanity.py b/tests/integration/defs/perf/test_perf_sanity.py index 0234e06dba7c..058f76c4f7ea 100644 --- a/tests/integration/defs/perf/test_perf_sanity.py +++ b/tests/integration/defs/perf/test_perf_sanity.py @@ -35,6 +35,7 @@ from ..conftest import get_llm_root, llm_models_root from ._model_paths import MODEL_PATH_DICT as _MODEL_PATH_DICT_BASE +from .disagg_server_config import build_disagg_server_config from .perf_regression_utils import process_and_upload_test_results # Sanity-side path differs from test_perf for this key; preserve historical value. @@ -1032,6 +1033,7 @@ def __init__( model_name: str, hardware: dict, server_env_var: str, + server_config_extra: Optional[Dict] = None, ): self.name = name self.disagg_serving_type = disagg_serving_type @@ -1042,6 +1044,7 @@ def __init__( self.model_name = model_name self.hardware = hardware self.server_env_var = server_env_var + self.server_config_extra = dict(server_config_extra or {}) self.num_ctx_servers = hardware.get("num_ctx_servers", 0) self.num_gen_servers = hardware.get("num_gen_servers", 0) @@ -1231,19 +1234,18 @@ def _generate_disagg_server_config(self, server_idx: int) -> str: # where another process on the same node grabs the port. disagg_server_port = get_free_port() - server_config = { - "hostname": self.hostname, - "port": disagg_server_port, - "backend": "pytorch", - "context_servers": { - "num_instances": self.num_ctx_servers, - "urls": ctx_hostnames, - }, - "generation_servers": { - "num_instances": self.num_gen_servers, - "urls": gen_hostnames, - }, - } + server_config_extra = {} + if server_idx < len(self.server_configs): + server_config_extra = self.server_configs[server_idx][2].server_config_extra + server_config = build_disagg_server_config( + hostname=self.hostname, + port=disagg_server_port, + num_ctx_servers=self.num_ctx_servers, + num_gen_servers=self.num_gen_servers, + ctx_hostnames=ctx_hostnames, + gen_hostnames=gen_hostnames, + server_config_extra=server_config_extra, + ) config_path = os.path.join(self.test_output_dir, f"server_config.{server_idx}.yaml") with open(config_path, "w") as f: yaml.dump(server_config, f) @@ -1827,6 +1829,7 @@ def _parse_disagg_config_file(self, config_file_path: str, config_file: str): model_name=model_name, hardware=hardware, server_env_var=server_env_var, + server_config_extra=config.get("server_config_extra", {}), ) # server_configs is a list with one element (tuple of ctx, gen, disagg config) diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_perf_sanity_ctx1_node1_gpu4_gen1_node2_gpu8.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_perf_sanity_ctx1_node1_gpu4_gen1_node2_gpu8.yml index 14aa2d0a4b24..3db665fb7efe 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_perf_sanity_ctx1_node1_gpu4_gen1_node2_gpu8.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_perf_sanity_ctx1_node1_gpu4_gen1_node2_gpu8.yml @@ -33,4 +33,4 @@ l0_gb300_multi_nodes_perf_sanity_ctx1_node1_gpu4_gen1_node2_gpu8: tests: # deepseek-r1-fp4 128k8k - perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL] TIMEOUT (180) - - perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL] TIMEOUT (180) + - perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL] TIMEOUT (220) diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 133a34e15300..051fb69970da 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -76,6 +76,11 @@ l0_h100: - unittest/disaggregated/test_extractor.py - unittest/disaggregated/test_peer.py - unittest/disaggregated/test_bounce.py + - unittest/disaggregated/test_transfer_quiescence.py + - unittest/disaggregated/test_async_consensus.py + # Launches a four-rank CPU protocol exercise with mpirun. This belongs in + # the MPI-enabled single-GPU stage; it does not require four physical GPUs. + - unittest/disaggregated/test_async_consensus_mpi.py - unittest/disaggregated/region/test_block.py - unittest/disaggregated/region/test_aux.py - unittest/disaggregated/region/test_page.py diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL.yaml index 92e9e36230a3..2021aef6671e 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL.yaml @@ -17,7 +17,7 @@ slurm: benchmark: mode: e2e use_nv_sa_benchmark: false - multi_round: 3 + multi_round: 1 benchmark_ratio: 0.0 streaming: true concurrency_list: '256' @@ -28,6 +28,8 @@ hardware: gpus_per_node: 4 num_ctx_servers: 1 num_gen_servers: 1 +server_config_extra: + schedule_style: generation_first environment: container_mount: container_image: @@ -37,6 +39,8 @@ environment: work_dir: worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=1 TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes + ctx_worker_env_var: TRTLLM_PYTHON_TRANSCEIVER_ASYNC_CTX_TERMINAL_CONSENSUS=0 TRTLLM_PYTHON_TRANSCEIVER_ASYNC_CTX_PEER_READY_CONSENSUS=0 + TRTLLM_PYTHON_TRANSCEIVER_CONTEXT_ACTIVATION_DIGEST=1 server_env_var: TRTLLM_SERVER_DISABLE_GC=1 profiling: nsys_on: false @@ -64,6 +68,7 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 131104 backend: NIXL + transceiver_runtime: PYTHON kv_transfer_timeout_ms: 600000 disable_overlap_scheduler: true speculative_config: &id001 @@ -91,6 +96,7 @@ worker_config: cache_transceiver_config: max_tokens_in_buffer: 131104 backend: NIXL + transceiver_runtime: PYTHON kv_transfer_timeout_ms: 600000 disable_overlap_scheduler: true speculative_config: *id001 diff --git a/tests/unittest/_torch/executor/test_async_transfer_manager.py b/tests/unittest/_torch/executor/test_async_transfer_manager.py index 1f2f9013d903..dad847640ef1 100644 --- a/tests/unittest/_torch/executor/test_async_transfer_manager.py +++ b/tests/unittest/_torch/executor/test_async_transfer_manager.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -15,7 +15,7 @@ from unittest.mock import MagicMock -from tensorrt_llm._torch.pyexecutor.py_executor import AsyncTransferManager +from tensorrt_llm._torch.pyexecutor.py_executor import AsyncTransferManager, AsyncTransferProvider from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType from tensorrt_llm.bindings import LlmRequestState @@ -112,6 +112,25 @@ def test_start_transfer_multiple_transfers_same_request(): kv_cache_manager.unpin_blocks_by_id.assert_called_once() +def test_end_transfer_retires_only_the_selected_provider(): + kv_cache_manager = MagicMock() + kv_cache_manager.store_blocks_for_reuse.return_value = 100 + resource_manager = create_mock_resource_manager(kv_cache_manager=kv_cache_manager) + manager = AsyncTransferManager(resource_manager) + request = create_mock_request(42) + + manager.start_transfer(request, AsyncTransferProvider.TRANSCEIVER) + manager.start_transfer(request, AsyncTransferProvider.CONNECTOR) + + assert not manager.end_transfer(request, AsyncTransferProvider.TRANSCEIVER) + assert not manager.has_transfer(request, AsyncTransferProvider.TRANSCEIVER) + assert manager.has_transfer(request, AsyncTransferProvider.CONNECTOR) + kv_cache_manager.unpin_blocks_by_id.assert_not_called() + + assert manager.end_transfer(request, AsyncTransferProvider.CONNECTOR) + kv_cache_manager.unpin_blocks_by_id.assert_called_once_with(100) + + def test_transfer_without_storing_blocks(): """Test starting a transfer with should_store_blocks=False.""" kv_cache_manager = MagicMock() diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py index 3d9a3f83dcc3..67291db312ed 100644 --- a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -14,6 +14,7 @@ # limitations under the License. import sys +import threading from types import SimpleNamespace from unittest.mock import Mock, call @@ -23,7 +24,12 @@ from tensorrt_llm._torch.pyexecutor import py_executor as executor_module from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import BindKvCacheTransceiver from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.pyexecutor.py_executor import ( + AsyncTransferManager, + AsyncTransferProvider, + PyExecutor, +) +from tensorrt_llm.bindings.executor import FinishReason from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig @@ -190,8 +196,10 @@ def test_flag_unset_context_timeout_preserves_legacy_cleanup(): request.py_kv_transfer_start_time = 1.0 request.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS executor = object.__new__(PyExecutor) + executor.canceled_req_ids = [] executor.kv_cache_transceiver = Mock() executor.kv_cache_transceiver.check_context_transfer_status.return_value = ([], []) + executor.kv_cache_transceiver.take_context_cancelled_request_ids.return_value = [] executor.kv_cache_transceiver.cancel_request.return_value = True executor.async_transfer_manager = Mock() executor.async_transfer_manager.requests_in_transfer.return_value = { @@ -208,7 +216,9 @@ def test_flag_unset_context_timeout_preserves_legacy_cleanup(): executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) assert request.py_kv_transfer_start_time is None assert request.state == LlmRequestState.DISAGG_CONTEXT_COMPLETE - executor._end_transfer_and_maybe_terminate.assert_called_once_with(request) + executor._end_transfer_and_maybe_terminate.assert_called_once_with( + request, AsyncTransferProvider.TRANSCEIVER + ) assert request.py_request_id not in executor._disagg_timed_out_ctx_cancelled_ids @@ -217,8 +227,10 @@ def test_enabled_context_timeout_defers_cleanup_until_cpp_terminal_state(monkeyp request.py_kv_transfer_start_time = 1.0 request.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS executor = object.__new__(PyExecutor) + executor.canceled_req_ids = [] executor.kv_cache_transceiver = Mock() executor.kv_cache_transceiver.check_context_transfer_status.return_value = ([], []) + executor.kv_cache_transceiver.take_context_cancelled_request_ids.return_value = [] executor.kv_cache_transceiver.cancel_request.return_value = True executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True executor.async_transfer_manager = Mock() @@ -253,7 +265,9 @@ def test_context_transfer_error_keeps_request_active_until_all_owners_release(): PyExecutor._end_transfer_and_maybe_terminate(executor, request) - executor.async_transfer_manager.end_transfer.assert_called_once_with(request) + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, AsyncTransferProvider.GENERIC + ) assert executor.active_requests == [request] executor._terminate_request.assert_not_called() @@ -301,6 +315,8 @@ def test_user_cancel_waits_for_context_transfer_owners(monkeypatch): executor.async_transfer_manager.requests_in_transfer.return_value = { request.py_request_id: request } + executor.async_transfer_manager.has_transfer.return_value = True + executor.async_transfer_manager.end_transfer.return_value = False monkeypatch.setattr(executor_module, "is_disagg_inflight_cancel_enabled", lambda: True) PyExecutor._handle_canceled_requests(executor) @@ -317,6 +333,575 @@ def test_user_cancel_waits_for_context_transfer_owners(monkeypatch): request.finish_by_reason.assert_called_once() +def test_user_cancel_routes_waiting_request_through_transceiver_owner(): + request = SimpleNamespace( + state=LlmRequestState.DISAGG_GENERATION_INIT, + py_request_id=7, + is_context_only_request=False, + ) + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.owns_request.return_value = True + executor.kv_cache_transceiver.cancel_request.return_value = False + + assert not PyExecutor._try_cancel_request(executor, request) + executor.kv_cache_transceiver.owns_request.assert_called_once_with(request) + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + + +def test_user_cancel_frees_unowned_request_without_transceiver_call(): + request = SimpleNamespace( + state=LlmRequestState.DISAGG_GENERATION_INIT, + py_request_id=7, + is_context_only_request=False, + ) + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.owns_request.return_value = False + + assert PyExecutor._try_cancel_request(executor, request) + executor.kv_cache_transceiver.cancel_request.assert_not_called() + + +def test_user_cancel_notifies_transceiver_owner_before_async_manager_release(): + request = SimpleNamespace( + state=LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER, + py_request_id=7, + is_context_only_request=True, + ) + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.owns_request.return_value = True + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = { + request.py_request_id: request + } + executor.async_transfer_manager.has_transfer.return_value = True + # A connector leg remains after the transceiver leg is acknowledged. + executor.async_transfer_manager.end_transfer.return_value = False + + assert not PyExecutor._try_cancel_request(executor, request) + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, AsyncTransferProvider.TRANSCEIVER + ) + assert request.py_request_id in executor._disagg_acknowledged_ctx_cancel_legs + + # Retrying must not decrement the connector leg a second time. + assert not PyExecutor._try_cancel_request(executor, request) + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, AsyncTransferProvider.TRANSCEIVER + ) + + executor.async_transfer_manager.requests_in_transfer.return_value = {} + executor.kv_cache_transceiver.owns_request.return_value = False + assert PyExecutor._try_cancel_request(executor, request) + assert request.py_request_id not in executor._disagg_acknowledged_ctx_cancel_legs + + +def test_user_cancel_does_not_retire_connector_after_transceiver_completion( + monkeypatch, +): + request = SimpleNamespace( + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + py_request_id=17, + is_context_only_request=True, + py_kv_transfer_start_time=1.0, + py_kv_transfer_timed_out=False, + ) + resource_manager = Mock() + resource_manager.resource_managers = {} + transfer_manager = AsyncTransferManager(resource_manager, should_store_blocks=False) + transfer_manager.start_transfer(request, AsyncTransferProvider.TRANSCEIVER) + transfer_manager.start_transfer(request, AsyncTransferProvider.CONNECTOR) + + executor = object.__new__(PyExecutor) + executor.async_transfer_manager = transfer_manager + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.owns_request.return_value = False + # A no-session cancellation historically reported success and retired the + # anonymous counter's remaining connector leg. + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.kv_cache_transceiver.kv_transfer_timeout_ms = 1 + executor.active_requests = [] + executor.force_terminate_ctx_for_partial_reuse = False + executor._terminate_request = Mock() + executor._is_disagg_inflight_cancel_active = Mock(return_value=True) + + PyExecutor._end_transfer_and_maybe_terminate( + executor, request, AsyncTransferProvider.TRANSCEIVER + ) + + assert not transfer_manager.has_transfer(request, AsyncTransferProvider.TRANSCEIVER) + assert transfer_manager.has_transfer(request, AsyncTransferProvider.CONNECTOR) + assert request.py_request_id in transfer_manager.requests_in_transfer() + + monkeypatch.setattr(executor_module.time, "monotonic", lambda: 10.0) + PyExecutor._check_kv_transfer_timeout(executor) + assert not request.py_kv_transfer_timed_out + + assert not PyExecutor._try_cancel_request(executor, request) + + executor.kv_cache_transceiver.cancel_request.assert_not_called() + executor._is_disagg_inflight_cancel_active.assert_not_called() + assert transfer_manager.has_transfer(request, AsyncTransferProvider.CONNECTOR) + assert request.py_request_id in transfer_manager.requests_in_transfer() + resource_manager.free_resources.assert_not_called() + executor._terminate_request.assert_not_called() + + +def test_duplicate_transceiver_completion_does_not_touch_connector_leg(): + request = SimpleNamespace( + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + py_request_id=18, + is_child=False, + is_context_only_request=True, + py_kv_transfer_timed_out=False, + ) + resource_manager = Mock() + resource_manager.resource_managers = {} + transfer_manager = AsyncTransferManager(resource_manager, should_store_blocks=False) + transfer_manager.start_transfer(request, AsyncTransferProvider.CONNECTOR) + + executor = object.__new__(PyExecutor) + executor.async_transfer_manager = transfer_manager + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.check_context_transfer_status.return_value = ( + [request.py_request_id], + [], + ) + executor.kv_cache_transceiver.take_context_cancelled_request_ids.return_value = [] + executor.canceled_req_ids = [] + executor._disagg_timed_out_ctx_cancelled_ids = set() + executor._end_transfer_and_maybe_terminate = Mock() + executor._check_cache_transfer_errors = Mock() + + PyExecutor._check_disagg_ctx_cache_transfer_status(executor, 0) + + executor._end_transfer_and_maybe_terminate.assert_not_called() + assert transfer_manager.has_transfer(request, AsyncTransferProvider.CONNECTOR) + assert request.py_request_id in transfer_manager.requests_in_transfer() + + +def test_peer_cancel_preserves_connector_only_leg_until_its_callback(): + request = SimpleNamespace( + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + py_request_id=19, + is_child=False, + is_context_only_request=True, + py_kv_transfer_timed_out=False, + ) + resource_manager = Mock() + resource_manager.resource_managers = {} + transfer_manager = AsyncTransferManager(resource_manager, should_store_blocks=False) + transfer_manager.start_transfer(request, AsyncTransferProvider.CONNECTOR) + + executor = object.__new__(PyExecutor) + executor.async_transfer_manager = transfer_manager + executor.active_requests = [request] + executor.canceled_req_ids = [] + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.check_context_transfer_status.return_value = ([], []) + executor.kv_cache_transceiver.take_context_cancelled_request_ids.return_value = [ + request.py_request_id + ] + executor._disagg_peer_cancelled_ctx_ids = set() + executor._disagg_timed_out_ctx_cancelled_ids = set() + executor._check_cache_transfer_errors = Mock() + + PyExecutor._check_disagg_ctx_cache_transfer_status(executor, 0) + + assert request.state == LlmRequestState.DISAGG_TRANS_ERROR + executor.kv_cache_transceiver.cancel_request.assert_not_called() + assert transfer_manager.has_transfer(request, AsyncTransferProvider.CONNECTOR) + assert request.py_request_id in transfer_manager.requests_in_transfer() + assert executor._disagg_peer_cancelled_ctx_ids == set() + + +@pytest.mark.parametrize( + "connector_final_state,is_child", + [ + (LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, False), + (LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, True), + (LlmRequestState.DISAGG_TRANS_ERROR, False), + (LlmRequestState.DISAGG_TRANS_ERROR, True), + ], + ids=[ + "connector-complete-request", + "connector-complete-child", + "connector-failed-request", + "connector-failed-child", + ], +) +def test_connector_callback_preserves_pending_user_cancel(connector_final_state, is_child): + request_id = 17 + parent_request_id = 117 + cancel_id = parent_request_id if is_child else request_id + response = SimpleNamespace(result=SimpleNamespace()) + request = SimpleNamespace( + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + py_request_id=request_id, + parent_request_id=parent_request_id, + is_child=is_child, + is_context_only_request=True, + is_dummy_request=False, + is_attention_dp_dummy=False, + py_kv_transfer_timed_out=False, + py_decoding_iter=2, + py_draft_tokens=[], + py_num_accepted_draft_tokens=0, + py_per_pos_drafted=[], + py_per_pos_accepted=[], + return_perf_metrics=False, + cached_tokens=0, + is_finished=False, + is_disagg_context_complete_state=False, + is_disagg_context_transmission_state=True, + is_generation_only_request=Mock(return_value=False), + create_response=Mock(return_value=response), + ) + + def finish_by_reason(reason): + assert reason == FinishReason.CANCELLED + request.state = LlmRequestState.GENERATION_COMPLETE + request.is_finished = True + request.is_disagg_context_complete_state = False + request.is_disagg_context_transmission_state = False + + request.finish_by_reason = Mock(side_effect=finish_by_reason) + + resource_manager = Mock() + resource_manager.resource_managers = {} + executor = object.__new__(PyExecutor) + executor.resource_manager = resource_manager + executor.async_transfer_manager = AsyncTransferManager( + resource_manager, should_store_blocks=False + ) + # One transceiver leg and one connector leg own the request. + executor.async_transfer_manager.start_transfer(request, AsyncTransferProvider.TRANSCEIVER) + executor.async_transfer_manager.start_transfer(request, AsyncTransferProvider.CONNECTOR) + executor.active_requests = [request] + executor.canceled_req_ids = [cancel_id] + executor.waiting_queue = Mock() + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.owns_request.return_value = True + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.kv_connector_manager = Mock() + executor.kv_connector_manager.get_finished.return_value = [request] + executor._disagg_acknowledged_ctx_cancel_legs = set() + executor._disagg_pp_termination_handler = None + executor._prefetched_request_ids = set() + executor._disagg_timed_out_ctx_cancelled_ids = set() + executor._disagg_timed_out_gen_cancelled_ids = set() + executor.result_wait_queues = {request_id: Mock()} + executor.gather_all_responses = False + executor.dist = SimpleNamespace(rank=0, world_size=1) + executor.perf_manager = Mock() + executor.perf_manager.get_timestamp.return_value = 0 + executor.iter_counter = 1 + executor.stream_interval = 1 + executor.force_terminate_ctx_for_partial_reuse = False + executor.enable_attention_dp = False + executor._pending_timed_out_requests = [] + executor._pending_transfer_responses = [] + executor._enqueue_responses = Mock() + executor._maybe_attach_ctx_usage = Mock() + executor._is_disagg_inflight_cancel_active = Mock(return_value=True) + + # User cancellation acknowledges only the transceiver leg. The connector + # still owns the request, so cancellation remains pending. + PyExecutor._handle_canceled_requests(executor) + + assert executor.canceled_req_ids == [cancel_id] + assert request_id in executor._disagg_acknowledged_ctx_cancel_legs + assert request_id in executor.async_transfer_manager.requests_in_transfer() + request.finish_by_reason.assert_not_called() + + # The real connector completion callback retires only its own leg. It must + # not create or enqueue a successful context response, even if that leg + # reports a failure while cancellation is pending. + request.state = connector_final_state + PyExecutor._kv_connector_terminate_requests(executor) + + assert request in executor.active_requests + assert request_id not in executor.async_transfer_manager.requests_in_transfer() + request.create_response.assert_not_called() + executor._enqueue_responses.assert_not_called() + assert executor._pending_transfer_responses == [] + resource_manager.free_resources.assert_not_called() + + if connector_final_state == LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS: + # AsyncTransferManager normally publishes CONTEXT_COMPLETE on its last + # leg. Reproduce the supported-mode race where transceiver state + # publication leaves the request's visible state stale. + request.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + + # The acknowledged transceiver leg plus an empty transfer manager is + # sufficient proof of quiescence even when inflight mode is active and the + # request state is stale. The ordinary response path then emits CANCELLED + # and frees the request exactly once. + PyExecutor._handle_canceled_requests(executor) + + assert executor.canceled_req_ids == [] + request.finish_by_reason.assert_called_once_with(FinishReason.CANCELLED) + executor.kv_cache_transceiver.owns_request.assert_called_once_with(request) + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + executor._is_disagg_inflight_cancel_active.assert_not_called() + assert request_id not in executor._disagg_acknowledged_ctx_cancel_legs + PyExecutor._handle_responses(executor) + + request.create_response.assert_called_once_with(False, 0) + executor._enqueue_responses.assert_called_once_with([(request_id, response)]) + assert executor.active_requests == [] + resource_manager.free_resources.assert_called_once_with(request) + assert request_id not in executor._disagg_acknowledged_ctx_cancel_legs + assert request_id not in executor.result_wait_queues + + +@pytest.mark.parametrize("status_mode", ["async", "legacy-default-off"]) +def test_peer_context_cancel_retires_transceiver_leg_as_error(status_mode): + request_id = 27 + request = SimpleNamespace( + state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS, + py_request_id=request_id, + parent_request_id=127, + is_child=False, + is_context_only_request=True, + is_dummy_request=False, + py_kv_transfer_timed_out=False, + create_response=Mock(), + ) + resource_manager = Mock() + resource_manager.resource_managers = {} + executor = object.__new__(PyExecutor) + executor.resource_manager = resource_manager + executor.async_transfer_manager = AsyncTransferManager( + resource_manager, should_store_blocks=False + ) + executor.async_transfer_manager.start_transfer(request, AsyncTransferProvider.TRANSCEIVER) + executor.active_requests = [request] + # This is a peer/global decision, not a locally queued user cancellation. + executor.canceled_req_ids = [] + executor.kv_cache_transceiver = Mock(name=f"{status_mode}_transceiver") + executor.kv_cache_transceiver.check_context_transfer_status.return_value = ([], []) + executor.kv_cache_transceiver.take_context_cancelled_request_ids.return_value = [request_id] + executor.kv_cache_transceiver.cancel_request.return_value = True + executor._disagg_peer_cancelled_ctx_ids = set() + executor._disagg_acknowledged_ctx_cancel_legs = set() + executor._disagg_timed_out_ctx_cancelled_ids = set() + executor._disagg_timed_out_gen_cancelled_ids = set() + executor._disagg_pp_termination_handler = None + executor._prefetched_request_ids = set() + executor.result_wait_queues = {request_id: Mock()} + executor.gather_all_responses = False + executor.enable_attention_dp = False + executor.dist = SimpleNamespace(rank=0, world_size=1) + observed_error_states = [] + + def handle_transfer_error(error_msg, requests, charge_budget): + assert error_msg == "Error in kv cache transfer for context requests" + assert charge_budget is False + assert requests == [request] + observed_error_states.append(request.state) + executor.active_requests.remove(request) + PyExecutor._terminate_request(executor, request) + + executor._handle_errors = Mock(side_effect=handle_transfer_error) + + PyExecutor._check_disagg_ctx_cache_transfer_status(executor, 0) + + assert observed_error_states == [LlmRequestState.DISAGG_TRANS_ERROR] + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + request.create_response.assert_not_called() + assert request_id not in executor.async_transfer_manager.requests_in_transfer() + assert executor.active_requests == [] + resource_manager.free_resources.assert_called_once_with(request) + assert executor._disagg_peer_cancelled_ctx_ids == set() + assert request_id not in executor.result_wait_queues + + +def test_user_cancel_acknowledges_default_off_transceiver_leg(monkeypatch): + request = SimpleNamespace( + state=LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER, + py_request_id=8, + is_context_only_request=True, + ) + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.owns_request.return_value = True + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = { + request.py_request_id: request + } + executor.async_transfer_manager.has_transfer.return_value = True + executor.async_transfer_manager.end_transfer.return_value = True + monkeypatch.setattr(executor_module, "is_disagg_inflight_cancel_enabled", lambda: False) + + assert PyExecutor._try_cancel_request(executor, request) + + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + executor.async_transfer_manager.end_transfer.assert_called_once_with( + request, AsyncTransferProvider.TRANSCEIVER + ) + assert request.py_request_id not in executor._disagg_acknowledged_ctx_cancel_legs + + +def test_known_cancel_withdraws_before_readiness_progression(): + request = SimpleNamespace( + py_request_id=9, + is_child=False, + is_context_only_request=True, + py_disaggregated_params=SimpleNamespace( + schedule_style=executor_module.DisaggScheduleStyle.GENERATION_FIRST + ), + ) + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.active_requests = [request] + executor.canceled_req_ids = [request.py_request_id] + call_order = Mock() + call_order.attach_mock( + executor.kv_cache_transceiver.exclude_context_requests_from_readiness, + "exclude", + ) + call_order.attach_mock( + executor.kv_cache_transceiver.prepare_context_requests, + "prepare", + ) + + PyExecutor._check_disagg_ctx_schedulable_status(executor, [request]) + + assert call_order.mock_calls == [call.exclude([request]), call.prepare([])] + + +def test_shutdown_bounds_transceiver_retries_before_registered_memory_release( + monkeypatch, +): + executor = object.__new__(PyExecutor) + executor.executor_request_queue = Mock() + executor.shutdown_event = Mock() + executor.hang_detector = Mock() + executor.hang_detector.detected.return_value = False + executor.worker_thread = Mock() + executor.dist = SimpleNamespace(pp_size=1) + executor._shutdown_sleep_wakeup_listeners = Mock() + executor.worker_started = True + model_engine = Mock() + executor.model_engine = model_engine + executor.draft_model_engine = None + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.shutdown.side_effect = RuntimeError("consensus shutdown failed") + resource_manager = Mock() + executor.resource_manager = SimpleNamespace(resource_managers={"kv": resource_manager}) + executor.virtual_memory_pools = {0: object()} + sampler = object.__new__(executor_module.AsyncWorkerMixin) + sampler._enable_async_worker = True + sampler.async_worker_stop = Mock() + executor.sampler = sampler + dwdp_exit = Mock() + executor.dwdp_manager = SimpleNamespace(__exit__=dwdp_exit) + monkeypatch.setattr(executor_module.torch.cuda, "is_available", lambda: False) + + with pytest.raises(RuntimeError, match="consensus shutdown failed"): + PyExecutor.shutdown(executor) + + assert executor.kv_cache_transceiver.shutdown.call_count == 3 + resource_manager.shutdown.assert_not_called() + model_engine._release_cuda_graphs.assert_called_once_with() + sampler.async_worker_stop.assert_called_once_with() + dwdp_exit.assert_called_once_with(None, None, None) + assert executor.dwdp_manager is None + assert executor.virtual_memory_pools == {} + assert not hasattr(executor, "model_engine") + + +def test_shutdown_retries_transient_transceiver_failure_before_resource_release( + monkeypatch, +): + executor = object.__new__(PyExecutor) + executor.executor_request_queue = Mock() + executor.shutdown_event = Mock() + executor.hang_detector = Mock() + executor.hang_detector.detected.return_value = False + executor.worker_thread = Mock() + executor.dist = SimpleNamespace(pp_size=1) + executor._shutdown_sleep_wakeup_listeners = Mock() + executor.worker_started = True + executor.model_engine = Mock() + executor.draft_model_engine = None + executor.kv_cache_transceiver = Mock() + call_order = [] + + def first_shutdown_attempt(): + call_order.append("transceiver-1") + raise RuntimeError("peer consensus shutdown timed out") + + def second_shutdown_attempt(): + call_order.append("transceiver-2") + + shutdown_attempts = iter([first_shutdown_attempt, second_shutdown_attempt]) + executor.kv_cache_transceiver.shutdown.side_effect = lambda: next(shutdown_attempts)() + resource_manager = Mock() + resource_manager.shutdown.side_effect = lambda: call_order.append("resource") + executor.resource_manager = SimpleNamespace(resource_managers={"kv": resource_manager}) + executor.virtual_memory_pools = None + executor.sampler = object() + executor.dwdp_manager = None + monkeypatch.setattr(executor_module.torch.cuda, "is_available", lambda: False) + + with pytest.raises(RuntimeError, match="peer consensus shutdown timed out"): + PyExecutor.shutdown(executor) + + assert call_order == ["transceiver-1", "transceiver-2", "resource"] + assert executor.kv_cache_transceiver.shutdown.call_count == 2 + resource_manager.shutdown.assert_called_once_with() + + +def test_shutdown_waits_for_deferred_transfer_teardown_before_resource_release( + monkeypatch, +): + executor = object.__new__(PyExecutor) + executor.executor_request_queue = Mock() + executor.shutdown_event = Mock() + executor.hang_detector = Mock() + executor.hang_detector.detected.return_value = False + executor.worker_thread = Mock() + executor.dist = SimpleNamespace(pp_size=1) + executor._shutdown_sleep_wakeup_listeners = Mock() + executor.worker_started = True + executor.model_engine = Mock() + executor.draft_model_engine = None + executor.kv_cache_transceiver = Mock() + completion = threading.Event() + executor.kv_cache_transceiver.shutdown.side_effect = [completion, None] + resource_manager = Mock() + + def release_registered_memory() -> None: + assert completion.is_set() + assert executor.kv_cache_transceiver.shutdown.call_count == 2 + + resource_manager.shutdown.side_effect = release_registered_memory + executor.resource_manager = SimpleNamespace(resource_managers={"kv": resource_manager}) + executor.virtual_memory_pools = None + executor.sampler = object() + executor.dwdp_manager = None + monkeypatch.setattr(executor_module.torch.cuda, "is_available", lambda: False) + + timer = threading.Timer(0.01, completion.set) + timer.start() + try: + PyExecutor.shutdown(executor) + finally: + timer.join() + + assert completion.is_set() + assert executor.kv_cache_transceiver.shutdown.call_count == 2 + resource_manager.shutdown.assert_called_once_with() + + def test_flag_unset_generation_driver_skips_cancel_pipeline(): executor = object.__new__(PyExecutor) executor.kv_cache_transceiver = Mock() diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 101edffbe155..a33002c6926f 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -16,6 +16,7 @@ import threading import time import types +from types import SimpleNamespace from unittest.mock import MagicMock, Mock import pytest @@ -33,6 +34,9 @@ ScheduledRequests, SerializableSchedulerOutput, ) +from tensorrt_llm.bindings.executor import FinishReason, RequestType +from tensorrt_llm.llmapi import DisaggScheduleStyle +from tensorrt_llm.llmapi.llm_args import WaitingQueuePolicy class MockPyExecutor: @@ -857,6 +861,9 @@ def test_pp_schedule_applies_gate_before_serializing(self): assert fitting == [] assert num_fitting == 0 assert wait_for_progress + executor.kv_cache_transceiver.activate_context_requests_for_schedule.assert_called_once_with( + [] + ) def test_pp_schedule_restores_propagated_gate_decision(self): executor = object.__new__(PyExecutor) @@ -869,6 +876,7 @@ def test_pp_schedule_restores_propagated_gate_decision(self): cp_size=1, ) executor.enable_attention_dp = False + executor.kv_cache_transceiver = None executor.active_requests = [ _make_disagg_transfer_request(1, 32, in_progress=True), _make_disagg_transfer_request(2, 32), @@ -893,6 +901,260 @@ def test_pp_schedule_restores_propagated_gate_decision(self): assert wait_for_progress +class TestGenerationFirstPreActiveRequests: + @staticmethod + def _request(request_id: int, *, priority: float = 0.0): + return SimpleNamespace( + request_id=request_id, + priority=priority, + request_type=RequestType.REQUEST_TYPE_CONTEXT_ONLY, + py_disaggregated_params=SimpleNamespace( + schedule_style=DisaggScheduleStyle.GENERATION_FIRST + ), + ) + + @classmethod + def _executor( + cls, + request_ids=(), + *, + max_num_active_requests: int = 4, + waiting_queue_policy=WaitingQueuePolicy.FCFS, + priorities=None, + ): + executor = object.__new__(PyExecutor) + executor.active_requests = [] + executor.max_num_active_requests = max_num_active_requests + executor._waiting_queue_policy = waiting_queue_policy + executor._gen_first_pre_active_requests = {} + executor._gen_first_pre_active_order = {} + executor._next_gen_first_pre_active_order = 0 + priorities = priorities or {} + for insertion_order, request_id in enumerate(request_ids): + request = cls._request(request_id, priority=priorities.get(request_id, 0.0)) + executor._gen_first_pre_active_requests[request_id] = request + executor._gen_first_pre_active_order[request_id] = insertion_order + executor._next_gen_first_pre_active_order += 1 + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.supports_pre_active_context_requests.return_value = True + return executor + + def test_default_off_leaves_context_request_in_waiting_queue(self): + executor = self._executor() + executor.kv_cache_transceiver.supports_pre_active_context_requests.return_value = False + queue = FCFSWaitingQueue() + item = RequestQueueItem(id=1, request=self._request(1)) + queue.add_request(item) + + assert PyExecutor._take_pre_active_context_items(executor, queue) == [] + assert list(queue) == [item] + + def test_context_first_request_stays_on_legacy_path(self): + executor = self._executor() + queue = FCFSWaitingQueue() + request = self._request(1) + request.py_disaggregated_params.schedule_style = DisaggScheduleStyle.CONTEXT_FIRST + item = RequestQueueItem(id=1, request=request) + queue.add_request(item) + + assert PyExecutor._take_pre_active_context_items(executor, queue) == [] + assert list(queue) == [item] + + def test_pre_active_extraction_does_not_consume_compute_capacity(self): + executor = self._executor(max_num_active_requests=1) + executor.active_requests = [self._request(99)] + queue = FCFSWaitingQueue() + items = [ + RequestQueueItem(id=request_id, request=self._request(request_id)) + for request_id in (1, 2, 3) + ] + queue.add_requests(items) + + assert PyExecutor._take_pre_active_context_items(executor, queue) == items + assert list(queue) == [] + assert executor.active_requests[0].request_id == 99 + + def test_rank_local_capacity_cannot_diverge_pre_active_cohort(self): + executors = [ + self._executor(max_num_active_requests=4), + self._executor(max_num_active_requests=4), + ] + executors[0].active_requests = [ + self._request(request_id) for request_id in (91, 92, 93, 94) + ] + queues = [] + for _ in executors: + queue = FCFSWaitingQueue() + queue.add_requests( + [ + RequestQueueItem(id=request_id, request=self._request(request_id)) + for request_id in (1, 2, 3) + ] + ) + queues.append(queue) + + cohorts = [ + [item.id for item in PyExecutor._take_pre_active_context_items(executor, queue)] + for executor, queue in zip(executors, queues) + ] + + assert cohorts == [[1, 2, 3], [1, 2, 3]] + assert all(not queue for queue in queues) + assert len(executors[0].active_requests) == 4 + assert executors[1].active_requests == [] + + def test_cancelled_qualified_item_still_uses_no_resource_lane(self): + executor = self._executor() + executor.canceled_req_ids = [1] + queue = FCFSWaitingQueue() + item = RequestQueueItem(id=1, request=self._request(1)) + queue.add_request(item) + + assert PyExecutor._take_pre_active_context_items(executor, queue) == [item] + assert list(queue) == [] + + def test_cancel_before_activation_is_excluded_before_readiness(self): + executor = self._executor() + request = Mock(request_id=9, py_request_id=9, is_child=False) + executor.canceled_req_ids = [9] + executor.waiting_queue = FCFSWaitingQueue() + executor._fetch_new_requests = Mock(return_value=([], [request])) + executor._validate_request = Mock() + + assert PyExecutor._fetch_and_activate_new_requests(executor) == [request] + + assert executor.active_requests == [] + assert executor._gen_first_pre_active_requests == {9: request} + executor.kv_cache_transceiver.exclude_context_requests_from_readiness.assert_called_once_with( + [request] + ) + executor.kv_cache_transceiver.prepare_context_requests.assert_not_called() + + def test_multi_return_request_uses_legacy_capacity_accounted_fallback(self): + executor = self._executor() + queue = FCFSWaitingQueue() + item = RequestQueueItem(id=1, request=self._request(1), child_req_ids=[101, 102]) + queue.add_request(item) + + assert PyExecutor._take_pre_active_context_items(executor, queue) == [] + assert list(queue) == [item] + + def test_selection_bypasses_unready_head_and_respects_capacity(self): + executor = self._executor((1, 2, 3), max_num_active_requests=2) + executor.active_requests = [self._request(99)] + executor.kv_cache_transceiver.is_context_request_ready_for_activation.side_effect = ( + lambda request: request.request_id in {2, 3} + ) + + assert PyExecutor._select_ready_pre_active_context_request_ids(executor) == [2] + + def test_priority_selection_is_canonical_across_ready_requests(self): + executor = self._executor( + (1, 2, 3), + max_num_active_requests=2, + waiting_queue_policy=WaitingQueuePolicy.PRIORITY, + priorities={1: 1.0, 2: 3.0, 3: 2.0}, + ) + executor.kv_cache_transceiver.is_context_request_ready_for_activation.return_value = True + + assert PyExecutor._select_ready_pre_active_context_request_ids(executor) == [2, 3] + + def test_exact_activation_ids_materialize_identically_on_follower(self): + scheduling_rank = self._executor((1, 2, 3), max_num_active_requests=2) + scheduling_rank.kv_cache_transceiver.is_context_request_ready_for_activation.side_effect = ( + lambda request: request.request_id in {2, 3} + ) + activation_ids = PyExecutor._activate_ready_pre_active_context_requests(scheduling_rank) + serialized = SerializableSchedulerOutput.from_scheduler_result( + ScheduledRequests(), + [], + 0, + activated_context_request_ids=activation_ids, + ) + + follower = self._executor((1, 2, 3), max_num_active_requests=2) + follower.kv_cache_transceiver.is_context_request_ready_for_activation.side_effect = ( + lambda request: request.request_id in {2, 3} + ) + PyExecutor._activate_pre_active_context_requests( + follower, serialized.activated_context_request_ids + ) + + assert activation_ids == [2, 3] + assert [request.request_id for request in scheduling_rank.active_requests] == [2, 3] + assert [request.request_id for request in follower.active_requests] == [2, 3] + assert list(follower._gen_first_pre_active_requests) == [1] + + def test_follower_missing_readiness_state_fails_before_mutation(self): + executor = self._executor((1, 2), max_num_active_requests=2) + executor.kv_cache_transceiver.is_context_request_ready_for_activation.side_effect = ( + lambda request: request.request_id == 1 + ) + + with pytest.raises(RuntimeError, match="without matching readiness"): + PyExecutor._activate_pre_active_context_requests(executor, [1, 2]) + + assert executor.active_requests == [] + assert list(executor._gen_first_pre_active_requests) == [1, 2] + executor.kv_cache_transceiver.activate_context_requests_for_schedule.assert_not_called() + + def test_idle_pre_active_cancellation_drains_without_model_resources(self): + executor = self._executor() + request = Mock( + request_id=7, + py_request_id=7, + is_child=False, + py_kv_transfer_timed_out=True, + py_decoding_iter=0, + py_draft_tokens=None, + cached_tokens=0, + ) + response = SimpleNamespace(result=SimpleNamespace(cached_tokens=None)) + request.create_response.return_value = response + executor._gen_first_pre_active_requests = {7: request} + executor._gen_first_pre_active_order = {7: 0} + executor.canceled_req_ids = [] + executor.is_shutdown = True + executor._try_cancel_request = Mock(return_value=True) + executor.dist = SimpleNamespace(rank=0) + executor.gather_all_responses = False + executor._maybe_attach_ctx_usage = Mock() + executor._enqueue_responses = Mock() + executor._prefetched_request_ids = {7} + executor._disagg_timed_out_ctx_cancelled_ids = {7} + executor._disagg_timed_out_gen_cancelled_ids = {7} + executor._disagg_acknowledged_ctx_cancel_legs = {7} + executor._disagg_peer_cancelled_ctx_ids = {7} + executor.result_wait_queues = {7: object()} + executor.resource_manager = Mock() + + PyExecutor._handle_pre_active_canceled_requests(executor) + + request.finish_by_reason.assert_called_once_with(FinishReason.CANCELLED) + executor._enqueue_responses.assert_called_once_with([(7, response)]) + executor.resource_manager.free_resources.assert_not_called() + assert executor._gen_first_pre_active_requests == {} + assert executor._gen_first_pre_active_order == {} + assert executor.canceled_req_ids == [] + assert 7 not in executor.result_wait_queues + + def test_generic_active_cancellation_preserves_pre_active_retry(self): + executor = self._executor() + request = Mock(request_id=8, py_request_id=8, is_child=False) + executor._gen_first_pre_active_requests = {8: request} + executor._gen_first_pre_active_order = {8: 0} + executor.active_requests = [] + executor.canceled_req_ids = [8] + executor.waiting_queue = FCFSWaitingQueue() + executor._try_cancel_request = Mock(return_value=True) + + PyExecutor._handle_canceled_requests(executor) + + executor._try_cancel_request.assert_not_called() + assert executor.canceled_req_ids == [8] + assert executor._gen_first_pre_active_requests == {8: request} + + def test_nonzero_pp_rank_prepares_snapshot_points_before_local_schedule( monkeypatch, ): diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py index c2a9a56a87bb..1a64dba7dbce 100644 --- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py +++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py @@ -88,7 +88,7 @@ def start_worker(self): class _DummyKvCacheCreator: """Mock KV cache creator that builds a dummy KV cache manager with reuse settings.""" - def __init__(self, **kwargs): + def __init__(self, *, estimate=False, **kwargs): """Initialize with KV cache configuration. Args: @@ -97,10 +97,19 @@ def __init__(self, **kwargs): self._max_seq_len = kwargs["max_seq_len"] self._kv_cache_config = kwargs["kv_cache_config"] self._execution_stream = kwargs["execution_stream"] + self._estimate = estimate def try_prepare_estimation(self): - """Skip estimation phase (no-op).""" - return False + """Return whether the mocked two-phase estimation path is active.""" + return self._estimate + + def configure_kv_cache_capacity(self, _py_executor): + """Accept the memory-accounting executor in the estimation path.""" + + @staticmethod + def teardown_managers(resources): + """Remove the temporary manager before final allocation.""" + resources.pop(ResourceManagerType.KV_CACHE_MANAGER, None) def build_managers(self, resources, estimating_kv_cache): """Build KV cache manager with reuse configuration. @@ -210,6 +219,8 @@ def _run_create_py_executor( enable_flash_mla=False, model_max_seq_len=128, enable_chunked_prefill=False, + estimate_kv_cache=False, + return_publish_flags=False, ): """Execute create_py_executor with mocked dependencies and return MLA runtime flags. @@ -225,6 +236,8 @@ def _run_create_py_executor( enable_flash_mla: Whether to emulate the FlashMLA block-size override. model_max_seq_len: Effective sequence length reported by the model engine. enable_chunked_prefill: Whether to request MLA chunked prefill support. + estimate_kv_cache: Whether to exercise two-phase KV cache estimation. + return_publish_flags: Whether to return endpoint-publication flags. Returns: Tuple of (kv_cache_reuse_flag, runtime_cache_reuse_flag, @@ -265,7 +278,11 @@ def _run_create_py_executor( monkeypatch.setattr(py_executor_creator, "is_mla", lambda _: True) monkeypatch.setattr(py_executor_creator, "is_hybrid_linear", lambda _: False) monkeypatch.setattr(py_executor_creator, "get_sm_version", lambda: sm_version) - monkeypatch.setattr(py_executor_creator, "KvCacheCreator", _DummyKvCacheCreator) + monkeypatch.setattr( + py_executor_creator, + "KvCacheCreator", + lambda **kwargs: _DummyKvCacheCreator(estimate=estimate_kv_cache, **kwargs), + ) monkeypatch.setattr(py_executor_creator.torch.cuda, "mem_get_info", lambda: (2 << 30, 4 << 30)) monkeypatch.setattr(py_executor_creator.torch.cuda, "empty_cache", lambda: None) @@ -287,7 +304,10 @@ def _create_model_engine(**kwargs): monkeypatch.setattr(py_executor_creator, "PyTorchModelEngine", _create_model_engine) + publish_flags = [] + def _create_py_executor_instance(**kwargs): + publish_flags.append(kwargs["publish_disaggregated_params"]) return _DummyPyExecutor( resources=kwargs["resources"], model_engine=kwargs["model_engine"], @@ -307,11 +327,28 @@ def _create_py_executor_instance(**kwargs): ResourceManagerType.KV_CACHE_MANAGER ) - return ( + result = ( kv_cache_manager.enable_block_reuse, py_executor.model_engine.attn_runtime_features.cache_reuse, py_executor.model_engine.attn_runtime_features.chunked_prefill, ) + return publish_flags if return_publish_flags else result + + +def test_kv_estimation_keeps_transceiver_but_suppresses_temporary_endpoint( + monkeypatch, +): + publish_flags = _run_create_py_executor( + monkeypatch, + sm_version=90, + kv_cache_quant_algo=QuantAlgo.NO_QUANT, + estimate_kv_cache=True, + return_publish_flags=True, + ) + + # Both executor lifetimes are constructed, preserving transceiver memory + # accounting. Only the final serving lifetime may publish its endpoints. + assert publish_flags == [False, True] def test_mla_unsupported_sm_fallback_syncs_cache_reuse(monkeypatch): diff --git a/tests/unittest/disaggregated/async_consensus_mpi_worker.py b/tests/unittest/disaggregated/async_consensus_mpi_worker.py new file mode 100644 index 000000000000..2cea51b66ad3 --- /dev/null +++ b/tests/unittest/disaggregated/async_consensus_mpi_worker.py @@ -0,0 +1,312 @@ +# 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. + +"""Four-rank CPU exercise for the asynchronous Python consensus protocol.""" + +from __future__ import annotations + +import json +import sys +import time +import traceback + +from mpi4py import MPI + +from tensorrt_llm._torch.disaggregation.async_consensus import ( + AsyncConsensusCoordinator, + ConsensusEvent, + ConsensusEventKind, + ConsensusOutcome, + ConsensusPhase, + MpiConsensusTransport, + _MessageKind, + _Packet, +) + +_WORLD_SIZE = 4 +_STEP_TIMEOUT_S = 10.0 + + +def _wait_for_event( + coordinator: AsyncConsensusCoordinator, + expected: set[ConsensusEventKind], + *, + reject: set[ConsensusEventKind] | None = None, +) -> ConsensusEvent: + deadline = time.monotonic() + _STEP_TIMEOUT_S + observed: list[str] = [] + while time.monotonic() < deadline: + for event in coordinator.poll(): + observed.append(event.kind.name) + if reject is not None and event.kind in reject: + raise AssertionError( + f"rank {coordinator.rank} observed rejected event {event}; history={observed}" + ) + if event.kind in expected: + return event + time.sleep(0.001) + raise AssertionError( + f"rank {coordinator.rank} timed out waiting for " + f"{sorted(kind.name for kind in expected)}; history={observed}" + ) + + +def _drain_pending(transport: MpiConsensusTransport) -> None: + deadline = time.monotonic() + _STEP_TIMEOUT_S + while transport.pending_send_count: + transport.progress() + if time.monotonic() >= deadline: + raise AssertionError( + f"rank {transport.rank} could not drain {transport.pending_send_count} sends" + ) + time.sleep(0.001) + + +def _run() -> None: + world = MPI.COMM_WORLD + rank = world.Get_rank() + if world.Get_size() != _WORLD_SIZE: + raise AssertionError(f"expected {_WORLD_SIZE} MPI ranks, got {world.Get_size()}") + + # A message on COMM_WORLD must remain untouched by the duplicated + # consensus communicator even though both use tag zero. + world_request = None + if rank == 0: + world_request = world.isend("world-communicator-sentinel", dest=3, tag=0) + + transport = MpiConsensusTransport(range(_WORLD_SIZE), max_pending_sends=256) + coordinator = AsyncConsensusCoordinator( + transport, + max_messages_per_poll=8, + max_open_rounds=64, + round_timeout_s=5.0, + ) + + # Terminal votes arrive in a different order on each rank. Failure must be + # reduced once, then delivered as the same authoritative outcome. + time.sleep(0.005 * ((_WORLD_SIZE - rank) % _WORLD_SIZE)) + local_outcome = ConsensusOutcome.FAILED if rank == 1 else ConsensusOutcome.COMPLETED + coordinator.publish_terminal(1001, local_outcome) + terminal = _wait_for_event( + coordinator, + {ConsensusEventKind.TERMINAL_COMMIT}, + ) + if terminal.outcome != ConsensusOutcome.FAILED: + raise AssertionError(f"rank {rank} observed terminal outcome {terminal.outcome}") + world.Barrier() + _drain_pending(transport) + world.Barrier() + + # Force a coordinator fan-out below its required three credits. The first + # attempt must emit nothing and retain a retryable action; restoring the + # capacity must then commit exactly once on every rank. + if rank == coordinator.coordinator_rank: + transport._max_pending_sends = 2 + world.Barrier() + coordinator.publish_terminal(1004, ConsensusOutcome.COMPLETED) + if rank == coordinator.coordinator_rank: + deadline = time.monotonic() + _STEP_TIMEOUT_S + while not coordinator._coordinator_actions: + if coordinator.poll(): + raise AssertionError("capacity-rejected fan-out emitted a local event") + if time.monotonic() >= deadline: + raise AssertionError("timed out reaching the capacity-rejected fan-out") + time.sleep(0.001) + if transport.pending_send_count: + raise AssertionError("capacity-rejected fan-out issued a partial MPI send") + transport._max_pending_sends = 256 + capacity_terminal = _wait_for_event( + coordinator, + {ConsensusEventKind.TERMINAL_COMMIT}, + ) + if capacity_terminal.outcome != ConsensusOutcome.COMPLETED: + raise AssertionError( + f"rank {rank} observed capacity retry outcome {capacity_terminal.outcome}" + ) + world.Barrier() + + # Exercise PREPARE racing a pre-ACK withdrawal. No scheduling rank may be + # released, and epoch reuse remains prohibited through rollback finalize. + coordinator.publish_ready(1002, epoch=0) + if rank == 2 and not coordinator.withdraw_ready(1002, epoch=0): + raise AssertionError("pre-ACK readiness withdrawal was rejected") + aborted = _wait_for_event( + coordinator, + {ConsensusEventKind.READY_ABORT}, + reject={ConsensusEventKind.READY_RELEASE, ConsensusEventKind.READY_COMPLETE}, + ) + if aborted.outcome != ConsensusOutcome.WITHDRAWN: + raise AssertionError(f"rank {rank} observed abort outcome {aborted.outcome}") + try: + coordinator.publish_ready(1002, epoch=1) + except RuntimeError as error: + if "before its prior epoch finalizes" not in str(error): + raise + else: + raise AssertionError("request ID was reused before readiness abort finalized") + coordinator.acknowledge_ready_abort(1002, epoch=0) + _wait_for_event(coordinator, {ConsensusEventKind.READY_ABORT_FINALIZE}) + world.Barrier() + + coordinator.publish_ready(1002, epoch=1) + _wait_for_event(coordinator, {ConsensusEventKind.READY_PREPARE}) + coordinator.acknowledge_ready(1002, epoch=1) + if rank == coordinator.scheduling_rank: + released = _wait_for_event(coordinator, {ConsensusEventKind.READY_RELEASE}) + if released.epoch != 1: + raise AssertionError(f"rank {rank} observed wrong ready epoch {released.epoch}") + elif rank == coordinator.coordinator_rank: + key = (ConsensusPhase.READY, 1002, 1) + deadline = time.monotonic() + _STEP_TIMEOUT_S + while key not in coordinator._ready_activation_required_acks: + coordinator.poll() + if time.monotonic() >= deadline: + raise AssertionError("timed out releasing the scheduling rank") + time.sleep(0.001) + world.Barrier() + # The barrier models delivery of rank zero's authoritative PP schedule. + coordinator.acknowledge_ready_activation(1002, epoch=1) + completed = _wait_for_event(coordinator, {ConsensusEventKind.READY_COMPLETE}) + if completed.epoch != 1: + raise AssertionError(f"rank {rank} observed wrong ready epoch {completed.epoch}") + world.Barrier() + + # A low-rank flood must not starve the two other senders when receiving one + # packet per poll. These packets exercise only the transport and are + # intentionally not passed to the coordinator state machine. + if rank < 3: + send_count = 8 if rank == 0 else 1 + for sequence in range(send_count): + transport.send( + _Packet( + _MessageKind.CLOSE, + ConsensusPhase.READY, + 2000 + sequence, + 0, + ConsensusOutcome.WITHDRAWN, + rank, + ), + 3, + ) + world.Barrier() + if rank == 3: + deadline = time.monotonic() + _STEP_TIMEOUT_S + while not all(transport._comm.Iprobe(source=source, tag=0) for source in (0, 1, 2)): + if time.monotonic() >= deadline: + raise AssertionError("timed out waiting for all fairness senders") + time.sleep(0.001) + first_sources = [transport.receive(1)[0].source for _ in range(3)] + if set(first_sources) != {0, 1, 2}: + raise AssertionError(f"rotating receive fairness failed: {first_sources}") + remaining = 7 + deadline = time.monotonic() + _STEP_TIMEOUT_S + while remaining: + packets = transport.receive(remaining) + remaining -= len(packets) + if time.monotonic() >= deadline: + raise AssertionError(f"timed out draining {remaining} fairness packets") + if not packets: + time.sleep(0.001) + world.Barrier() + _drain_pending(transport) + + if rank == 3: + sentinel = world.recv(source=0, tag=0) + if sentinel != "world-communicator-sentinel": + raise AssertionError(f"unexpected communicator sentinel: {sentinel}") + if rank == 0: + world_request.wait() + world.Barrier() + + # A deliberately incomplete round must diagnose and stop; a watchdog must + # never fabricate a terminal commit to make progress. + coordinator._round_timeout_s = 0.1 + if rank == 0: + coordinator.publish_terminal(1003, ConsensusOutcome.COMPLETED) + watchdog_seen = False + deadline = time.monotonic() + 1.0 + while time.monotonic() < deadline: + try: + events = coordinator.poll() + except RuntimeError as error: + if not any( + marker in str(error) + for marker in ( + "watchdog expired without a global decision", + "coordinated fail-stop", + ) + ): + raise + watchdog_seen = True + break + if any(event.kind == ConsensusEventKind.TERMINAL_COMMIT for event in events): + raise AssertionError("watchdog path invented a terminal consensus outcome") + time.sleep(0.001) + if not watchdog_seen: + raise AssertionError( + f"rank {rank} watchdog state mismatch: observed={watchdog_seen}, " + f"coordinator={coordinator.coordinator_rank}" + ) + world.Barrier() + + # Rank zero times out once before the other ranks enter shutdown, then + # retries the same close handshake. CLOSE must remain idempotent and the + # staggered communicator free must finish on every rank. + retried_shutdown = False + if rank == 0: + try: + coordinator.shutdown(0.02) + except RuntimeError as error: + if "shutdown acknowledgement" not in str(error): + raise + retried_shutdown = True + else: + raise AssertionError("rank zero shutdown unexpectedly completed before peers") + time.sleep(0.15) + else: + time.sleep(0.05 + 0.01 * rank) + coordinator.shutdown(10.0) + world.Barrier() + + retries = world.gather(retried_shutdown, root=0) + if rank == 0: + if retries != [True, False, False, False]: + raise AssertionError(f"unexpected shutdown retry flags: {retries}") + print( + "ASYNC_CONSENSUS_MPI_OK " + + json.dumps( + { + "communicator_isolation": True, + "fanout_backpressure_retry": True, + "epoch_reuse": True, + "receive_fairness": True, + "shutdown_retry": True, + "terminal_outcome": terminal.outcome.name, + "watchdog_fail_closed": True, + }, + sort_keys=True, + ), + flush=True, + ) + + +if __name__ == "__main__": + try: + _run() + except BaseException: + traceback.print_exc() + sys.stderr.flush() + MPI.COMM_WORLD.Abort(1) + raise diff --git a/tests/unittest/disaggregated/test_async_consensus.py b/tests/unittest/disaggregated/test_async_consensus.py new file mode 100644 index 000000000000..33ff03612d01 --- /dev/null +++ b/tests/unittest/disaggregated/test_async_consensus.py @@ -0,0 +1,1488 @@ +# 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. + +from __future__ import annotations + +import threading +import time +from collections import defaultdict, deque +from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest + +import tensorrt_llm._torch.disaggregation.async_consensus as async_consensus_module +from tensorrt_llm._torch.disaggregation.async_consensus import ( + AsyncConsensusCoordinator, + ConsensusEvent, + ConsensusEventKind, + ConsensusOutcome, + ConsensusPhase, + MpiConsensusTransport, + _ConsensusBackpressure, + _CoordinatorAction, + _MessageKind, + _Packet, + _PendingSend, +) + + +class _FakeNetwork: + def __init__(self, participants: Sequence[int]) -> None: + self.participants = tuple(participants) + self._queues: dict[int, deque[_Packet]] = {rank: deque() for rank in self.participants} + self._lock = threading.Lock() + + def send(self, packet: _Packet, destination: int) -> None: + with self._lock: + self._queues[destination].append(packet) + + def receive(self, rank: int, limit: int) -> list[_Packet]: + packets: list[_Packet] = [] + with self._lock: + queue = self._queues[rank] + while queue and len(packets) < limit: + packets.append(queue.popleft()) + return packets + + def queued(self, rank: int) -> int: + with self._lock: + return len(self._queues[rank]) + + +class _FakeTransport: + def __init__(self, network: _FakeNetwork, rank: int) -> None: + self.rank = rank + self.participants = network.participants + self._network = network + self.receive_limits: list[int] = [] + self.progress_count = 0 + self.close_timeouts: list[float] = [] + + def send(self, packet: _Packet, destination: int) -> None: + self._network.send(packet, destination) + + def send_many(self, messages: Sequence[tuple[_Packet, int]]) -> None: + for packet, destination in messages: + self.send(packet, destination) + + def progress(self) -> None: + self.progress_count += 1 + + def receive(self, limit: int) -> list[_Packet]: + self.receive_limits.append(limit) + return self._network.receive(self.rank, limit) + + @property + def pending_send_count(self) -> int: + return 0 + + def close(self, timeout_s: float) -> None: + self.close_timeouts.append(timeout_s) + + +class _AtomicCapacityTransport(_FakeTransport): + def __init__(self, network: _FakeNetwork, rank: int, capacity: int) -> None: + super().__init__(network, rank) + self.capacity = capacity + + def send_many(self, messages: Sequence[tuple[_Packet, int]]) -> None: + if len(messages) > self.capacity: + raise _ConsensusBackpressure( + f"send backpressure limit exceeded: batch={len(messages)}, limit={self.capacity}" + ) + super().send_many(messages) + + +class _FanoutFailureTransport(_FakeTransport): + """Reject one matching fan-out before delivering any packet.""" + + def __init__( + self, + network: _FakeNetwork, + rank: int, + failing_kind: _MessageKind, + *, + defer_fail_stop_once: bool = False, + ) -> None: + super().__init__(network, rank) + self._failing_kind = failing_kind + self._defer_fail_stop_once = defer_fail_stop_once + self._fail_stop_deferred = False + self.failed = False + self.fanout_kinds: list[_MessageKind] = [] + + def send_many(self, messages: Sequence[tuple[_Packet, int]]) -> None: + messages_tuple = tuple(messages) + if messages_tuple: + self.fanout_kinds.append(messages_tuple[0][0].kind) + if ( + messages_tuple + and messages_tuple[0][0].kind == _MessageKind.FAIL_STOP + and self._defer_fail_stop_once + and not self._fail_stop_deferred + ): + self._fail_stop_deferred = True + raise _ConsensusBackpressure("synthetic fail-stop backpressure") + if not self.failed and messages_tuple and messages_tuple[0][0].kind == self._failing_kind: + self.failed = True + raise RuntimeError(f"synthetic {self._failing_kind.name} fan-out failure") + super().send_many(messages_tuple) + + +class _ReceiveFailureTransport(_FakeTransport): + def __init__(self, network: _FakeNetwork, rank: int, diagnostic: str) -> None: + super().__init__(network, rank) + self._diagnostic = diagnostic + self._failed = False + + def receive(self, limit: int) -> list[_Packet]: + if not self._failed: + self._failed = True + raise RuntimeError(self._diagnostic) + return super().receive(limit) + + +class _FakeClock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +class _NoIterationDict(dict): + def __iter__(self): + raise AssertionError("hot-path state must not be globally scanned") + + +def _make_group( + size: int = 4, + *, + max_messages_per_poll: int = 256, + max_completed_epochs: int = 65_536, + max_open_rounds: int = 65_536, + round_timeout_s: float = 600.0, + ready_lease_timeout_s: float | None = None, + clock: _FakeClock | None = None, +) -> tuple[_FakeNetwork, list[_FakeTransport], list[AsyncConsensusCoordinator]]: + network = _FakeNetwork(range(size)) + transports = [_FakeTransport(network, rank) for rank in range(size)] + coordinators = [ + AsyncConsensusCoordinator( + transport, + max_messages_per_poll=max_messages_per_poll, + max_completed_epochs=max_completed_epochs, + max_open_rounds=max_open_rounds, + round_timeout_s=round_timeout_s, + ready_lease_timeout_s=ready_lease_timeout_s, + clock=clock if clock is not None else time.monotonic, + ) + for transport in transports + ] + return network, transports, coordinators + + +def _poll_rounds( + coordinators: Sequence[AsyncConsensusCoordinator], + rounds: int = 4, +) -> dict[int, list[ConsensusEvent]]: + events: dict[int, list[ConsensusEvent]] = defaultdict(list) + for _ in range(rounds): + for coordinator in coordinators: + events[coordinator.rank].extend(coordinator.poll()) + return events + + +def _terminal_events(events: Sequence[ConsensusEvent]) -> list[ConsensusEvent]: + return [event for event in events if event.kind == ConsensusEventKind.TERMINAL_COMMIT] + + +def _release_scheduler_activate_all_and_complete( + coordinators: Sequence[AsyncConsensusCoordinator], + request_id: int, + epoch: int = 0, +) -> dict[int, list[ConsensusEvent]]: + """Use the authoritative schedule token, then complete every lease.""" + scheduling_rank = coordinators[0].scheduling_rank + coordinator = coordinators[-1] + events: dict[int, list[ConsensusEvent]] = defaultdict(list) + + assert coordinator.poll() == [] + scheduler = next(instance for instance in coordinators if instance.rank == scheduling_rank) + release_events = scheduler.poll() + events[scheduler.rank].extend(release_events) + assert [event.kind for event in release_events] == [ConsensusEventKind.READY_RELEASE] + + # The test now simulates delivery of the exact rank-zero PP schedule. + for instance in coordinators: + instance.acknowledge_ready_activation(request_id, epoch) + completed = _poll_rounds(coordinators) + for rank, rank_events in completed.items(): + events[rank].extend(rank_events) + assert [event.kind for event in rank_events] == [ConsensusEventKind.READY_COMPLETE] + return events + + +def test_terminal_success_waits_for_every_staggered_vote() -> None: + _, _, coordinators = _make_group() + request_id = 101 + + for rank in (2, 0, 3): + coordinators[rank].publish_terminal(request_id, ConsensusOutcome.COMPLETED) + events = _poll_rounds(coordinators) + assert all(not _terminal_events(rank_events) for rank_events in events.values()) + + coordinators[1].publish_terminal(request_id, ConsensusOutcome.COMPLETED) + events = _poll_rounds(coordinators) + + for rank in range(4): + assert _terminal_events(events[rank]) == [ + ConsensusEvent( + ConsensusEventKind.TERMINAL_COMMIT, + request_id, + 0, + ConsensusOutcome.COMPLETED, + ) + ] + + +@pytest.mark.parametrize( + ("votes", "expected"), + [ + ( + [ + ConsensusOutcome.COMPLETED, + ConsensusOutcome.FAILED, + ConsensusOutcome.COMPLETED, + ConsensusOutcome.COMPLETED, + ], + ConsensusOutcome.FAILED, + ), + ( + [ + ConsensusOutcome.FAILED, + ConsensusOutcome.CANCELLED, + ConsensusOutcome.COMPLETED, + ConsensusOutcome.FAILED, + ], + ConsensusOutcome.CANCELLED, + ), + ], +) +def test_terminal_outcome_precedence( + votes: list[ConsensusOutcome], expected: ConsensusOutcome +) -> None: + _, _, coordinators = _make_group() + + for coordinator, outcome in zip(coordinators, votes): + coordinator.publish_terminal(102, outcome) + events = _poll_rounds(coordinators) + + for rank_events in events.values(): + assert [event.outcome for event in _terminal_events(rank_events)] == [expected] + + +def test_duplicate_vote_is_idempotent_and_changed_vote_is_rejected() -> None: + network, _, coordinators = _make_group() + request_id = 103 + follower = coordinators[0] + coordinator = coordinators[-1] + + follower.publish_terminal(request_id, ConsensusOutcome.COMPLETED) + follower.publish_terminal(request_id, ConsensusOutcome.COMPLETED) + assert network.queued(coordinator.rank) == 1 + with pytest.raises(RuntimeError, match="local consensus vote changed"): + follower.publish_terminal(request_id, ConsensusOutcome.FAILED) + + duplicate = _Packet( + _MessageKind.VOTE, + ConsensusPhase.TERMINAL, + request_id, + 0, + ConsensusOutcome.COMPLETED, + follower.rank, + ) + network.send(duplicate, coordinator.rank) + coordinator.poll() + + changed = _Packet( + _MessageKind.VOTE, + ConsensusPhase.TERMINAL, + request_id, + 0, + ConsensusOutcome.FAILED, + follower.rank, + ) + network.send(changed, coordinator.rank) + with pytest.raises(RuntimeError, match="participant 0 changed vote"): + coordinator.poll() + + +@pytest.mark.parametrize( + ("kind", "phase", "outcome"), + [ + (_MessageKind.READY_PREPARE, ConsensusPhase.READY, ConsensusOutcome.READY), + (_MessageKind.READY_RELEASE, ConsensusPhase.READY, ConsensusOutcome.READY), + (_MessageKind.READY_COMPLETE, ConsensusPhase.READY, ConsensusOutcome.READY), + (_MessageKind.READY_ABORT, ConsensusPhase.READY, ConsensusOutcome.WITHDRAWN), + ( + _MessageKind.READY_ABORT_FINALIZE, + ConsensusPhase.READY, + ConsensusOutcome.WITHDRAWN, + ), + ( + _MessageKind.TERMINAL_COMMIT, + ConsensusPhase.TERMINAL, + ConsensusOutcome.COMPLETED, + ), + (_MessageKind.CLOSE_ACK, ConsensusPhase.READY, ConsensusOutcome.WITHDRAWN), + ], +) +def test_authoritative_messages_must_come_from_coordinator( + kind: _MessageKind, + phase: ConsensusPhase, + outcome: ConsensusOutcome, +) -> None: + network, _, coordinators = _make_group() + follower = coordinators[0] + network.send(_Packet(kind, phase, 1031, 0, outcome, 1), follower.rank) + + with pytest.raises(RuntimeError, match=rf"message {kind.name} did not come from coordinator"): + follower.poll() + assert not follower._events + assert follower._fatal_key == (phase, 1031, 0) + assert follower._fatal_error is not None + assert network.queued(coordinators[-1].rank) == 1 + + +def test_packet_contract_rejects_wrong_ack_outcome() -> None: + network, _, coordinators = _make_group() + coordinator = coordinators[-1] + network.send( + _Packet( + _MessageKind.READY_ACK, + ConsensusPhase.READY, + 1032, + 0, + ConsensusOutcome.FAILED, + 0, + ), + coordinator.rank, + ) + + with pytest.raises(RuntimeError, match="message READY_ACK has outcome FAILED"): + coordinator.poll() + assert coordinator._fatal_key == (ConsensusPhase.READY, 1032, 0) + assert coordinator._fatal_error == "message READY_ACK has outcome FAILED, expected READY" + assert coordinator._fail_stop_propagated + for follower in coordinators[:-1]: + with pytest.raises(RuntimeError, match="coordinated fail-stop"): + follower.poll() + + +def test_protocol_version_mismatch_is_fatal_not_retryable_backpressure() -> None: + fields = _Packet( + _MessageKind.VOTE, + ConsensusPhase.TERMINAL, + 1033, + 0, + ConsensusOutcome.COMPLETED, + 0, + ).encode() + fields[0] = np.uint64(0) + + with pytest.raises(RuntimeError, match="unsupported consensus protocol version") as error: + _Packet.decode(fields) + assert type(error.value) is RuntimeError + + +def test_receive_decode_failure_preserves_first_error_and_notifies_all_ranks() -> None: + network = _FakeNetwork(range(4)) + diagnostic = "synthetic packet decode failure" + transports: list[_FakeTransport] = [ + _ReceiveFailureTransport(network, 0, diagnostic), + *[_FakeTransport(network, rank) for rank in range(1, 4)], + ] + coordinators = [AsyncConsensusCoordinator(transport) for transport in transports] + follower = coordinators[0] + + with pytest.raises(RuntimeError, match=diagnostic): + follower.poll() + assert follower._fatal_error == diagnostic + assert follower._fatal_key == (ConsensusPhase.TERMINAL, 0, 0) + assert not follower._priority_local_outbox + + # The decode failure's reserved notification reaches the coordinator, + # which fans out authoritative fail-stop to every participant. + with pytest.raises(RuntimeError, match="coordinated fail-stop"): + coordinators[-1].poll() + for instance in coordinators[1:-1]: + with pytest.raises(RuntimeError, match="coordinated fail-stop"): + instance.poll() + + # Further polling must report the first local diagnostic, not overwrite it + # with the coordinator's later generic fail-stop message. + with pytest.raises(RuntimeError, match=diagnostic): + follower.poll() + assert follower._fatal_error == diagnostic + + +@pytest.mark.parametrize( + "failing_kind", + [_MessageKind.READY_PREPARE, _MessageKind.TERMINAL_COMMIT], +) +def test_decision_fanout_hard_failure_is_replaced_by_fail_stop( + failing_kind: _MessageKind, +) -> None: + network = _FakeNetwork(range(4)) + transports: list[_FakeTransport] = [_FakeTransport(network, rank) for rank in range(3)] + coordinator_transport = _FanoutFailureTransport( + network, + 3, + failing_kind, + defer_fail_stop_once=True, + ) + transports.append(coordinator_transport) + coordinators = [AsyncConsensusCoordinator(transport) for transport in transports] + coordinator = coordinators[-1] + + for instance in coordinators: + if failing_kind == _MessageKind.READY_PREPARE: + instance.publish_ready(1034) + else: + instance.publish_terminal(1034, ConsensusOutcome.COMPLETED) + + diagnostic = f"synthetic {failing_kind.name} fan-out failure" + with pytest.raises(RuntimeError, match=diagnostic): + coordinator.poll() + assert coordinator._fatal_error == diagnostic + assert not coordinator._fail_stop_propagated + assert coordinator_transport.fanout_kinds == [ + failing_kind, + _MessageKind.FAIL_STOP, + ] + assert [action for action, _ in coordinator._coordinator_actions] == [ + _CoordinatorAction.FAIL_STOP + ] + assert not coordinator._events + + # A later poll retries only the reserved fail-stop action. The failed + # decision can never be retried. + with pytest.raises(RuntimeError, match=diagnostic): + coordinator.poll() + assert coordinator._fail_stop_propagated + assert not coordinator._coordinator_actions + assert coordinator_transport.fanout_kinds == [ + failing_kind, + _MessageKind.FAIL_STOP, + _MessageKind.FAIL_STOP, + ] + + # No participant can observe a decision from the rejected atomic fan-out. + expected_phase = ( + ConsensusPhase.READY + if failing_kind == _MessageKind.READY_PREPARE + else ConsensusPhase.TERMINAL + ) + for follower in coordinators[:-1]: + with pytest.raises(RuntimeError, match="coordinated fail-stop"): + follower.poll() + assert not follower._events + assert follower._fatal_key == (expected_phase, 1034, 0) + + # Subsequent progress reports the original diagnostic and can never retry + # the failed decision action. + with pytest.raises(RuntimeError, match=diagnostic): + coordinator.poll() + assert coordinator_transport.fanout_kinds == [ + failing_kind, + _MessageKind.FAIL_STOP, + _MessageKind.FAIL_STOP, + ] + + +def test_publication_and_withdrawal_absorb_send_backpressure() -> None: + network, transports, coordinators = _make_group() + follower = coordinators[0] + transport = transports[0] + original_send = transport.send + fail_next = True + + def fail_once(packet: _Packet, destination: int) -> None: + nonlocal fail_next + if fail_next: + fail_next = False + raise _ConsensusBackpressure("synthetic send backpressure") + original_send(packet, destination) + + transport.send = fail_once + terminal_key = (ConsensusPhase.TERMINAL, 104, 0) + follower.publish_terminal(104, ConsensusOutcome.COMPLETED) + assert network.queued(coordinators[-1].rank) == 0 + assert follower._local_votes[terminal_key] == ConsensusOutcome.COMPLETED + assert (_MessageKind.VOTE, terminal_key) in follower._local_outbox + follower.poll() + assert network.queued(coordinators[-1].rank) == 1 + assert follower._local_votes[terminal_key] == ConsensusOutcome.COMPLETED + assert not follower._local_outbox + + follower.publish_ready(105) + ready_key = (ConsensusPhase.READY, 105, 0) + fail_next = True + assert follower.withdraw_ready(105) is True + assert follower._local_votes[ready_key] == ConsensusOutcome.WITHDRAWN + assert (_MessageKind.WITHDRAW, ready_key) in follower._local_outbox + follower.poll() + assert not follower._local_outbox + + +def test_hard_send_failure_enters_local_fail_stop_and_queues_notification() -> None: + _, transports, coordinators = _make_group() + follower = coordinators[0] + transport = transports[follower.rank] + + def fail_send(_packet: _Packet, _destination: int) -> None: + raise RuntimeError("synthetic hard transport failure") + + transport.send = fail_send + with pytest.raises(RuntimeError, match="synthetic hard transport failure"): + follower.publish_terminal(1041, ConsensusOutcome.COMPLETED) + + key = (ConsensusPhase.TERMINAL, 1041, 0) + assert follower._fatal_key == key + assert not follower._local_outbox + assert list(follower._priority_local_outbox) == [(_MessageKind.FAIL_STOP, key)] + with pytest.raises(RuntimeError, match="after coordinated fail-stop starts"): + follower.publish_terminal(1042, ConsensusOutcome.COMPLETED) + + +def test_unsent_readiness_vote_is_coalesced_into_withdrawal() -> None: + network, transports, coordinators = _make_group() + follower = coordinators[0] + transport = transports[0] + original_send = transport.send + + def reject_send(_packet: _Packet, _destination: int) -> None: + raise _ConsensusBackpressure("synthetic send backpressure") + + transport.send = reject_send + follower.publish_ready(1051) + key = (ConsensusPhase.READY, 1051, 0) + assert list(follower._local_outbox) == [(_MessageKind.VOTE, key)] + + assert follower.withdraw_ready(1051) is True + assert list(follower._local_outbox) == [(_MessageKind.WITHDRAW, key)] + + transport.send = original_send + follower.poll() + queued = network.receive(coordinators[-1].rank, 2) + assert [packet.kind for packet in queued] == [_MessageKind.WITHDRAW] + + +def test_ready_abort_supersedes_unsent_normal_intent_for_same_round() -> None: + network, transports, coordinators = _make_group(max_open_rounds=1) + follower = coordinators[0] + transport = transports[follower.rank] + + def reject_send(_packet: _Packet, _destination: int) -> None: + raise _ConsensusBackpressure("synthetic send backpressure") + + transport.send = reject_send + follower.publish_ready(1052) + key = (ConsensusPhase.READY, 1052, 0) + assert list(follower._local_outbox) == [(_MessageKind.VOTE, key)] + + network.send( + _Packet( + _MessageKind.READY_ABORT, + ConsensusPhase.READY, + 1052, + 0, + ConsensusOutcome.WITHDRAWN, + follower.coordinator_rank, + ), + follower.rank, + ) + assert [event.kind for event in follower.poll()] == [ConsensusEventKind.READY_ABORT] + assert not follower._local_outbox + + # The reserved round credit is sufficient for the only remaining outbound + # obligation even though transport pressure persists. + follower.acknowledge_ready_abort(1052) + assert list(follower._local_outbox) == [(_MessageKind.READY_ABORT_ACK, key)] + + +def test_ready_ack_intents_survive_backpressure_after_prepare_event_is_consumed() -> None: + _, transports, coordinators = _make_group() + coordinator = coordinators[-1] + follower = coordinators[1] + + for instance in coordinators: + instance.publish_ready(1052) + coordinator.poll() + events_by_rank = {instance.rank: instance.poll() for instance in coordinators[:-1]} + coordinator_events = [ + ConsensusEvent( + ConsensusEventKind.READY_PREPARE, + 1052, + 0, + ConsensusOutcome.READY, + ) + ] + assert all(events == coordinator_events for events in events_by_rank.values()) + + original_send = transports[follower.rank].send + reject_next_ack = True + + def reject_ack_once(packet: _Packet, destination: int) -> None: + nonlocal reject_next_ack + if reject_next_ack and packet.kind == _MessageKind.READY_ACK: + reject_next_ack = False + raise _ConsensusBackpressure("synthetic ACK backpressure") + original_send(packet, destination) + + transports[follower.rank].send = reject_ack_once + follower.acknowledge_ready(1052) + key = (ConsensusPhase.READY, 1052, 0) + assert key in follower._local_ready_acknowledged + assert (_MessageKind.READY_ACK, key) in follower._local_outbox + + for instance in coordinators: + if instance is not follower: + instance.acknowledge_ready(1052) + assert coordinator.poll() == [] + + follower.poll() + _release_scheduler_activate_all_and_complete(coordinators, 1052) + + +def test_ready_activation_ack_survives_backpressure_before_completion() -> None: + _, transports, coordinators = _make_group() + request_id = 1053 + coordinator = coordinators[-1] + follower = coordinators[1] + scheduling_rank = coordinators[0] + + for instance in coordinators: + instance.publish_ready(request_id) + coordinator.poll() + for instance in coordinators[:-1]: + instance.poll() + for instance in coordinators: + instance.acknowledge_ready(request_id) + + assert coordinator.poll() == [] + assert [event.kind for event in scheduling_rank.poll()] == [ConsensusEventKind.READY_RELEASE] + + original_send = transports[follower.rank].send + reject_next_ack = True + + def reject_activation_ack_once(packet: _Packet, destination: int) -> None: + nonlocal reject_next_ack + if reject_next_ack and packet.kind == _MessageKind.READY_ACTIVATE_ACK: + reject_next_ack = False + raise _ConsensusBackpressure("synthetic activation-ACK backpressure") + original_send(packet, destination) + + transports[follower.rank].send = reject_activation_ack_once + follower.acknowledge_ready_activation(request_id) + key = (ConsensusPhase.READY, request_id, 0) + assert (_MessageKind.READY_ACTIVATE_ACK, key) in follower._local_outbox + + scheduling_rank.acknowledge_ready_activation(request_id) + coordinators[2].acknowledge_ready_activation(request_id) + coordinator.acknowledge_ready_activation(request_id) + assert coordinator.poll() == [] + + follower.poll() + assert [event.kind for event in coordinator.poll()] == [ConsensusEventKind.READY_COMPLETE] + assert [event.kind for event in scheduling_rank.poll()] == [ConsensusEventKind.READY_COMPLETE] + + +def test_terminal_fanout_reserves_capacity_before_any_rank_commits() -> None: + network = _FakeNetwork(range(4)) + transports: list[_FakeTransport] = [_FakeTransport(network, rank) for rank in range(3)] + coordinator_transport = _AtomicCapacityTransport(network, 3, capacity=2) + transports.append(coordinator_transport) + coordinators = [AsyncConsensusCoordinator(transport) for transport in transports] + + for instance in coordinators: + instance.publish_terminal(106, ConsensusOutcome.COMPLETED) + assert coordinators[-1].poll() == [] + + assert all(network.queued(rank) == 0 for rank in range(3)) + assert not coordinators[-1]._events + action = coordinators[-1]._coordinator_actions[0] + assert action[1] == (ConsensusPhase.TERMINAL, 106, 0) + + coordinator_transport.capacity = 3 + events = {3: coordinators[-1].poll()} + for instance in coordinators[:-1]: + events[instance.rank] = instance.poll() + assert all(len(_terminal_events(rank_events)) == 1 for rank_events in events.values()) + + +def test_readiness_prepare_ack_release_orders_scheduler_last() -> None: + _, _, coordinators = _make_group() + request_id = 201 + scheduling_rank = coordinators[0] + coordinator = coordinators[-1] + assert all(instance.scheduling_rank == scheduling_rank.rank for instance in coordinators) + + for instance in coordinators: + instance.publish_ready(request_id) + + coordinator_events = coordinator.poll() + assert coordinator_events == [ + ConsensusEvent( + ConsensusEventKind.READY_PREPARE, + request_id, + 0, + ConsensusOutcome.READY, + ) + ] + scheduling_events = scheduling_rank.poll() + rank_one_events = coordinators[1].poll() + rank_two_events = coordinators[2].poll() + assert [event.kind for event in scheduling_events] == [ConsensusEventKind.READY_PREPARE] + assert [event.kind for event in rank_one_events] == [ConsensusEventKind.READY_PREPARE] + assert [event.kind for event in rank_two_events] == [ConsensusEventKind.READY_PREPARE] + + scheduling_rank.acknowledge_ready(request_id) + coordinators[1].acknowledge_ready(request_id) + coordinator.acknowledge_ready(request_id) + assert coordinator.poll() == [] + + coordinators[2].acknowledge_ready(request_id) + events = _release_scheduler_activate_all_and_complete(coordinators, request_id) + for rank in range(4): + assert ConsensusEventKind.READY_COMPLETE in [event.kind for event in events[rank]] + assert ConsensusEventKind.READY_RELEASE in [ + event.kind for event in events[scheduling_rank.rank] + ] + + +def test_released_readiness_lease_does_not_use_protocol_idle_watchdog() -> None: + clock = _FakeClock() + _, _, coordinators = _make_group(round_timeout_s=5.0, clock=clock) + request_id = 2021 + coordinator = coordinators[-1] + scheduling_rank = coordinators[0] + + for instance in coordinators: + instance.publish_ready(request_id) + assert [event.kind for event in coordinator.poll()] == [ConsensusEventKind.READY_PREPARE] + for instance in coordinators[:-1]: + assert [event.kind for event in instance.poll()] == [ConsensusEventKind.READY_PREPARE] + for instance in coordinators: + instance.acknowledge_ready(request_id) + assert coordinator.poll() == [] + assert [event.kind for event in scheduling_rank.poll()] == [ConsensusEventKind.READY_RELEASE] + + # Rank zero may wait arbitrarily longer than the vote/ack watchdog for KV + # capacity. No consensus packet is expected during this healthy lease. + clock.advance(50.0) + assert all(instance.poll() == [] for instance in coordinators) + + for instance in coordinators: + instance.acknowledge_ready_activation(request_id) + events = _poll_rounds(coordinators) + assert all( + [event.kind for event in events[rank]] == [ConsensusEventKind.READY_COMPLETE] + for rank in range(len(coordinators)) + ) + + +def test_readiness_ack_before_prepare_is_rejected() -> None: + _, _, coordinators = _make_group() + follower = coordinators[1] + + follower.publish_ready(203) + with pytest.raises(RuntimeError, match="before READY_PREPARE"): + follower.acknowledge_ready(203) + + +def test_readiness_withdraw_aborts_epoch_and_tombstones_late_messages() -> None: + network, _, coordinators = _make_group() + request_id = 202 + coordinator = coordinators[-1] + + for instance in coordinators: + instance.publish_ready(request_id) + prepare_events = coordinator.poll() + assert [event.kind for event in prepare_events] == [ConsensusEventKind.READY_PREPARE] + + coordinators[1].withdraw_ready(request_id) + coordinator_events = coordinator.poll() + assert [event.kind for event in coordinator_events] == [ConsensusEventKind.READY_ABORT] + abort_events = {coordinator.rank: coordinator_events} + for instance in coordinators[:-1]: + abort_events[instance.rank] = instance.poll() + assert [event.kind for event in abort_events[0]] == [ + ConsensusEventKind.READY_PREPARE, + ConsensusEventKind.READY_ABORT, + ] + # READY_PREPARE was already in flight when rank 1 withdrew. The withdrawing + # rank must suppress that stale preparation, while another follower may + # legally observe PREPARE followed by ABORT. Neither path may release the + # scheduling rank. + assert [event.kind for event in abort_events[1]] == [ConsensusEventKind.READY_ABORT] + assert [event.kind for event in abort_events[2]] == [ + ConsensusEventKind.READY_PREPARE, + ConsensusEventKind.READY_ABORT, + ] + assert [event.kind for event in abort_events[coordinator.rank]] == [ + ConsensusEventKind.READY_ABORT + ] + + # READY_ABORT is only the rollback command. No rank may reuse the request + # ID until every rollback is applied and the coordinator finalizes it. + for instance in coordinators: + with pytest.raises(RuntimeError, match="before its prior epoch finalizes"): + instance.publish_ready(request_id, epoch=1) + for instance in (coordinators[0], coordinators[1], coordinator): + instance.acknowledge_ready_abort(request_id) + + # Finalization is itself a consensus point: one unapplied rollback keeps + # the request ID/epoch leased everywhere. + assert coordinator.poll() == [] + coordinators[2].acknowledge_ready_abort(request_id) + + coordinator_events = coordinator.poll() + assert [event.kind for event in coordinator_events] == [ConsensusEventKind.READY_ABORT_FINALIZE] + finalize_events = {coordinator.rank: coordinator_events} + for instance in coordinators[:-1]: + finalize_events[instance.rank] = instance.poll() + for rank_events in finalize_events.values(): + assert [event.kind for event in rank_events] == [ConsensusEventKind.READY_ABORT_FINALIZE] + + stale_vote = _Packet( + _MessageKind.VOTE, + ConsensusPhase.READY, + request_id, + 0, + ConsensusOutcome.READY, + 0, + ) + network.send(stale_vote, coordinator.rank) + assert coordinator.poll() == [] + with pytest.raises(RuntimeError, match="stale consensus epoch"): + coordinators[0].publish_ready(request_id, epoch=0) + + for instance in coordinators: + instance.publish_ready(request_id, epoch=1) + assert [event.kind for event in coordinator.poll()] == [ConsensusEventKind.READY_PREPARE] + + +def test_readiness_ack_is_an_irrevocable_lease() -> None: + network, _, coordinators = _make_group() + request_id = 205 + coordinator = coordinators[-1] + + for instance in coordinators: + instance.publish_ready(request_id) + coordinator.poll() + for instance in coordinators[:-1]: + instance.poll() + + coordinators[1].acknowledge_ready(request_id) + assert coordinators[1].withdraw_ready(request_id) is False + # Even a malformed late withdrawal cannot overturn a lease whose ACK was + # observed first from the same source. + network.send( + _Packet( + _MessageKind.WITHDRAW, + ConsensusPhase.READY, + request_id, + 0, + ConsensusOutcome.WITHDRAWN, + 1, + ), + coordinator.rank, + ) + coordinators[0].acknowledge_ready(request_id) + coordinators[2].acknowledge_ready(request_id) + coordinator.acknowledge_ready(request_id) + + _release_scheduler_activate_all_and_complete(coordinators, request_id) + + +def test_withdraw_before_ack_aborts_even_after_another_rank_leases() -> None: + _, _, coordinators = _make_group() + request_id = 206 + coordinator = coordinators[-1] + + for instance in coordinators: + instance.publish_ready(request_id) + coordinator.poll() + for instance in coordinators[:-1]: + instance.poll() + + coordinators[1].acknowledge_ready(request_id) + assert coordinators[2].withdraw_ready(request_id) is True + coordinator_events = coordinator.poll() + assert [event.kind for event in coordinator_events] == [ConsensusEventKind.READY_ABORT] + + events = {coordinator.rank: coordinator_events} + for instance in coordinators[:-1]: + events[instance.rank] = instance.poll() + assert all( + ConsensusEventKind.READY_RELEASE not in [event.kind for event in rank_events] + for rank_events in events.values() + ) + assert [event.kind for event in events[1]] == [ConsensusEventKind.READY_ABORT] + assert coordinators[1].withdraw_ready(request_id) is True + + for instance in coordinators: + instance.acknowledge_ready_abort(request_id) + assert [event.kind for event in coordinator.poll()] == [ConsensusEventKind.READY_ABORT_FINALIZE] + for instance in coordinators[:-1]: + assert [event.kind for event in instance.poll()] == [ + ConsensusEventKind.READY_ABORT_FINALIZE + ] + + +def test_ready_abort_ack_intent_survives_backpressure() -> None: + _, transports, coordinators = _make_group() + coordinator = coordinators[-1] + follower = coordinators[1] + request_id = 2061 + + for instance in coordinators: + instance.publish_ready(request_id) + coordinator.poll() + for instance in coordinators[:-1]: + instance.poll() + assert coordinators[2].withdraw_ready(request_id) is True + assert [event.kind for event in coordinator.poll()] == [ConsensusEventKind.READY_ABORT] + for instance in coordinators[:-1]: + assert ConsensusEventKind.READY_ABORT in [event.kind for event in instance.poll()] + + original_send = transports[follower.rank].send + reject_next_ack = True + + def reject_abort_ack_once(packet: _Packet, destination: int) -> None: + nonlocal reject_next_ack + if reject_next_ack and packet.kind == _MessageKind.READY_ABORT_ACK: + reject_next_ack = False + raise _ConsensusBackpressure("synthetic abort-ACK backpressure") + original_send(packet, destination) + + transports[follower.rank].send = reject_abort_ack_once + for instance in coordinators: + instance.acknowledge_ready_abort(request_id) + key = (ConsensusPhase.READY, request_id, 0) + assert key in follower._local_ready_abort_acknowledged + assert (_MessageKind.READY_ABORT_ACK, key) in follower._local_outbox + assert coordinator.poll() == [] + + follower.poll() + assert [event.kind for event in coordinator.poll()] == [ConsensusEventKind.READY_ABORT_FINALIZE] + for instance in coordinators[:-1]: + assert [event.kind for event in instance.poll()] == [ + ConsensusEventKind.READY_ABORT_FINALIZE + ] + + +def test_withdraw_without_local_vote_aborts_partial_readiness_round() -> None: + _, _, coordinators = _make_group() + request_id = 207 + coordinator = coordinators[-1] + + # One rank has observed peer metadata and voted. A second rank receives a + # cancellation before it becomes locally ready, while the remaining ranks + # have not entered the round at all. + coordinators[0].publish_ready(request_id) + assert coordinators[1].withdraw_ready(request_id) is True + + coordinator_events = coordinator.poll() + assert [event.kind for event in coordinator_events] == [ConsensusEventKind.READY_ABORT] + events = {coordinator.rank: coordinator_events} + for instance in coordinators[:-1]: + events[instance.rank] = instance.poll() + assert all( + [event.kind for event in rank_events] == [ConsensusEventKind.READY_ABORT] + for rank_events in events.values() + ) + + # Every participant rolls back and acknowledges before any rank can reuse + # the request ID at the next epoch. + for instance in coordinators: + with pytest.raises(RuntimeError, match="before its prior epoch finalizes"): + instance.publish_ready(request_id, epoch=1) + instance.acknowledge_ready_abort(request_id) + assert [event.kind for event in coordinator.poll()] == [ConsensusEventKind.READY_ABORT_FINALIZE] + for instance in coordinators[:-1]: + assert [event.kind for event in instance.poll()] == [ + ConsensusEventKind.READY_ABORT_FINALIZE + ] + for instance in coordinators: + instance.publish_ready(request_id, epoch=1) + + +def test_poll_processes_at_most_configured_messages() -> None: + network, transports, coordinators = _make_group(max_messages_per_poll=2) + coordinator = coordinators[-1] + + for request_id in range(3): + coordinators[0].publish_terminal(request_id, ConsensusOutcome.COMPLETED) + assert network.queued(coordinator.rank) == 3 + + assert coordinator.poll() == [] + assert network.queued(coordinator.rank) == 1 + assert transports[coordinator.rank].receive_limits == [2] + assert coordinator.poll() == [] + assert network.queued(coordinator.rank) == 0 + assert transports[coordinator.rank].receive_limits == [2, 2] + + +def test_newer_commit_purges_incomplete_older_epoch() -> None: + _, _, coordinators = _make_group() + coordinator = coordinators[-1] + request_id = 204 + + coordinators[0].publish_terminal(request_id, ConsensusOutcome.COMPLETED, epoch=0) + coordinator.poll() + + for instance in coordinators: + instance.publish_terminal(request_id, ConsensusOutcome.COMPLETED, epoch=1) + _poll_rounds(coordinators) + + assert (ConsensusPhase.TERMINAL, request_id, 0) not in coordinator._votes + + +def test_completion_and_ready_epoch_checks_use_direct_indexes() -> None: + _, _, coordinators = _make_group(max_open_rounds=512) + follower = coordinators[0] + for request_id in range(9000, 9256): + follower.publish_terminal(request_id, ConsensusOutcome.COMPLETED) + + follower._local_votes = _NoIterationDict(follower._local_votes) + completed_key = (ConsensusPhase.TERMINAL, 9000, 0) + follower._complete_local(completed_key) + assert completed_key not in follower._local_votes + assert len(follower._round_progress) == 255 + + follower.publish_ready(9300) + follower.publish_ready(9301) + with pytest.raises(RuntimeError, match="before its prior epoch finalizes"): + follower.publish_ready(9300, epoch=1) + + +def test_completed_epoch_tombstones_have_a_bounded_lru_window() -> None: + _, _, coordinators = _make_group(max_completed_epochs=2) + + for request_id in (401, 402, 403): + for instance in coordinators: + instance.publish_terminal(request_id, ConsensusOutcome.COMPLETED) + _poll_rounds(coordinators) + + for instance in coordinators: + assert len(instance._completed_epoch) == 2 + assert (ConsensusPhase.TERMINAL, 401) not in instance._completed_epoch + assert list(instance._completed_epoch) == [ + (ConsensusPhase.TERMINAL, 402), + (ConsensusPhase.TERMINAL, 403), + ] + + +def test_open_round_limit_rejects_new_work_without_mutating_state() -> None: + network, _, coordinators = _make_group(max_open_rounds=2) + follower = coordinators[0] + coordinator = coordinators[-1] + + follower.publish_terminal(501, ConsensusOutcome.COMPLETED) + follower.publish_terminal(502, ConsensusOutcome.COMPLETED) + with pytest.raises(RuntimeError, match="open-round limit exceeded"): + follower.publish_terminal(503, ConsensusOutcome.COMPLETED) + + assert network.queued(coordinator.rank) == 2 + assert (ConsensusPhase.TERMINAL, 503, 0) not in follower._local_votes + + +def test_coordinator_open_round_limit_bounds_untrusted_remote_votes() -> None: + network, _, coordinators = _make_group(max_open_rounds=2) + coordinator = coordinators[-1] + + for request_id in (511, 512, 513): + network.send( + _Packet( + _MessageKind.VOTE, + ConsensusPhase.TERMINAL, + request_id, + 0, + ConsensusOutcome.COMPLETED, + 0, + ), + coordinator.rank, + ) + + with pytest.raises(RuntimeError, match="operation=coordinator receive"): + coordinator.poll() + assert len(coordinator._round_progress) == 2 + assert (ConsensusPhase.TERMINAL, 513, 0) not in coordinator._votes + + +def test_round_watchdog_reports_missing_consensus_and_never_commits() -> None: + clock = _FakeClock() + _, _, coordinators = _make_group(round_timeout_s=5.0, clock=clock) + follower = coordinators[0] + + follower.publish_terminal(521, ConsensusOutcome.COMPLETED) + assert follower.poll() == [] + clock.advance(5.1) + + with pytest.raises( + RuntimeError, + match=r"watchdog expired.*phase=TERMINAL.*request_id=521.*missing_votes=\[0, 1, 2, 3\]", + ): + follower.poll() + assert follower._events == deque() + + +def test_round_watchdog_refreshes_on_progress_then_expires_when_idle() -> None: + clock = _FakeClock() + _, _, coordinators = _make_group(round_timeout_s=5.0, clock=clock) + coordinator = coordinators[-1] + + coordinators[0].publish_terminal(525, ConsensusOutcome.COMPLETED) + assert coordinator.poll() == [] + + # A second distinct vote is healthy protocol progress. It arrives near the + # original deadline and must grant the round a fresh idle window. + clock.advance(4.0) + coordinators[1].publish_terminal(525, ConsensusOutcome.COMPLETED) + assert coordinator.poll() == [] + + clock.advance(1.1) + assert coordinator.poll() == [] + + # No further vote or acknowledgement arrives, so the refreshed idle + # window still fails closed once it expires. + clock.advance(4.0) + with pytest.raises(RuntimeError, match=r"watchdog expired.*request_id=525"): + coordinator.poll() + + +def test_watchdog_propagates_coordinated_fail_stop_to_uninvolved_ranks() -> None: + clock = _FakeClock() + _, _, coordinators = _make_group(round_timeout_s=5.0, clock=clock) + coordinators[0].publish_terminal(522, ConsensusOutcome.COMPLETED) + clock.advance(5.1) + + with pytest.raises(RuntimeError, match="watchdog expired"): + coordinators[0].poll() + with pytest.raises(RuntimeError, match="coordinated fail-stop"): + coordinators[-1].poll() + for instance in coordinators[1:-1]: + with pytest.raises(RuntimeError, match="coordinated fail-stop"): + instance.poll() + + assert all(not instance._events for instance in coordinators) + with pytest.raises(RuntimeError, match="shutdown acknowledgement"): + coordinators[1].shutdown(0.001) + with ThreadPoolExecutor(max_workers=len(coordinators)) as executor: + futures = [executor.submit(instance.shutdown, 2.0) for instance in coordinators] + for future in futures: + future.result(timeout=3.0) + + +def test_coordinator_watchdog_fails_locally_when_fail_stop_fanout_is_backpressured() -> None: + clock = _FakeClock() + network = _FakeNetwork(range(4)) + transports: list[_FakeTransport] = [_FakeTransport(network, rank) for rank in range(3)] + coordinator_transport = _AtomicCapacityTransport(network, 3, capacity=0) + transports.append(coordinator_transport) + coordinators = [ + AsyncConsensusCoordinator(transport, round_timeout_s=1.0, clock=clock) + for transport in transports + ] + coordinator = coordinators[-1] + + for instance in coordinators: + instance.publish_terminal(523, ConsensusOutcome.COMPLETED) + assert coordinator.poll() == [] + clock.advance(1.1) + + with pytest.raises(RuntimeError, match=r"watchdog expired.*request_id=523"): + coordinator.poll() + assert coordinator._fatal_key == (ConsensusPhase.TERMINAL, 523, 0) + assert coordinator._fatal_error is not None + assert not coordinator._fail_stop_propagated + assert coordinator._coordinator_actions[0][0] == _CoordinatorAction.FAIL_STOP + assert all(network.queued(rank) == 0 for rank in range(3)) + + +def test_follower_watchdog_fails_locally_and_reserves_notification_under_backpressure() -> None: + clock = _FakeClock() + _, transports, coordinators = _make_group(round_timeout_s=1.0, clock=clock) + follower = coordinators[0] + + def reject_send(_packet: _Packet, _destination: int) -> None: + raise _ConsensusBackpressure("permanent synthetic backpressure") + + transports[follower.rank].send = reject_send + follower.publish_terminal(524, ConsensusOutcome.COMPLETED) + clock.advance(1.1) + + with pytest.raises(RuntimeError, match=r"watchdog expired.*request_id=524"): + follower.poll() + key = (ConsensusPhase.TERMINAL, 524, 0) + assert follower._fatal_key == key + assert not follower._local_outbox + assert list(follower._priority_local_outbox) == [(_MessageKind.FAIL_STOP, key)] + + +def test_ready_action_queue_does_not_scan_incomplete_rounds() -> None: + _, _, coordinators = _make_group(max_messages_per_poll=256) + coordinator = coordinators[-1] + + for request_id in range(100): + coordinators[0].publish_terminal(request_id, ConsensusOutcome.COMPLETED) + coordinator.poll() + assert not coordinator._coordinator_actions + + request_id = 601 + for instance in coordinators: + instance.publish_terminal(request_id, ConsensusOutcome.COMPLETED) + events = coordinator.poll() + assert _terminal_events(events) == [ + ConsensusEvent( + ConsensusEventKind.TERMINAL_COMMIT, + request_id, + 0, + ConsensusOutcome.COMPLETED, + ) + ] + + +class _ControllableRequest: + def __init__(self) -> None: + self.complete = False + self.test_count = 0 + + def Test(self) -> bool: + self.test_count += 1 + return self.complete + + +class _FakeMpiComm: + def __init__(self) -> None: + self.requests: list[_ControllableRequest] = [] + self.sent_destinations: list[int] = [] + self.receive_queues: dict[int, deque[_Packet]] = defaultdict(deque) + self.error_handler = None + self.freed = False + + def Get_size(self) -> int: + return 4 + + def Dup(self): + return self + + def Get_rank(self) -> int: + return 3 + + def Set_errhandler(self, error_handler) -> None: + self.error_handler = error_handler + + def Isend(self, _buffer, dest: int, tag: int) -> _ControllableRequest: + assert tag == 0 + request = _ControllableRequest() + self.requests.append(request) + self.sent_destinations.append(dest) + return request + + def Iprobe(self, source: int, tag: int) -> bool: + assert tag == 0 + return bool(self.receive_queues[source]) + + def Recv(self, fields, source: int, tag: int) -> None: + assert tag == 0 + fields[:] = self.receive_queues[source].popleft().encode() + + def Free(self) -> None: + self.freed = True + + +def _make_unit_mpi_transport( + comm: _FakeMpiComm, + *, + max_pending_sends: int = 4, + max_send_tests_per_progress: int = 256, +) -> MpiConsensusTransport: + transport = object.__new__(MpiConsensusTransport) + transport.participants = (0, 1, 2, 3) + transport._comm = comm + transport.rank = 3 + transport._pending = deque() + transport._max_pending_sends = max_pending_sends + transport._max_send_tests_per_progress = max_send_tests_per_progress + transport._receive_sources = (0, 1, 2) + transport._receive_cursor = 0 + transport._closed = False + return transport + + +def test_mpi_transport_makes_post_preflight_errors_process_fatal( + monkeypatch: pytest.MonkeyPatch, +) -> None: + comm = _FakeMpiComm() + fatal_handler = object() + monkeypatch.setattr( + async_consensus_module, + "MPI", + SimpleNamespace(ERRORS_ARE_FATAL=fatal_handler), + ) + monkeypatch.setattr(async_consensus_module, "mpi_comm", lambda: comm) + + transport = MpiConsensusTransport(range(4)) + + assert transport._comm is comm + assert comm.error_handler is fatal_handler + + +def test_mpi_transport_frees_duplicated_communicator_when_setup_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + comm = _FakeMpiComm() + comm.Set_errhandler = Mock(side_effect=RuntimeError("MPI setup failed")) + monkeypatch.setattr( + async_consensus_module, + "MPI", + SimpleNamespace(ERRORS_ARE_FATAL=object()), + ) + monkeypatch.setattr(async_consensus_module, "mpi_comm", lambda: comm) + + with pytest.raises(RuntimeError, match="MPI setup failed"): + MpiConsensusTransport(range(4)) + + assert comm.freed + + +def test_mpi_transport_applies_send_backpressure_before_allocating() -> None: + comm = _FakeMpiComm() + transport = _make_unit_mpi_transport(comm, max_pending_sends=1) + key = (ConsensusPhase.TERMINAL, 701, 0) + packet = _Packet( + _MessageKind.VOTE, + key[0], + key[1], + key[2], + ConsensusOutcome.COMPLETED, + transport.rank, + ) + + transport.send(packet, 0) + with pytest.raises( + _ConsensusBackpressure, + match=r"backpressure limit exceeded.*pending=1, projected=2, limit=1", + ): + transport.send(packet, 1) + assert comm.sent_destinations == [0] + assert transport.pending_send_count == 1 + + comm.requests[0].complete = True + transport.send(packet, 1) + assert comm.sent_destinations == [0, 1] + assert transport.pending_send_count == 1 + + +def test_mpi_transport_progress_checks_a_bounded_rotating_window() -> None: + comm = _FakeMpiComm() + transport = _make_unit_mpi_transport( + comm, + max_pending_sends=16, + max_send_tests_per_progress=3, + ) + packet = _Packet( + _MessageKind.VOTE, + ConsensusPhase.TERMINAL, + 702, + 0, + ConsensusOutcome.COMPLETED, + transport.rank, + ) + for destination in (0, 1, 2, 0, 1, 2): + # Bypass send() so its own bounded progress does not affect the test + # window being measured. + buffer = packet.encode() + transport._pending.append( + _PendingSend(buffer=buffer, request=comm.Isend(buffer, dest=destination, tag=0)) + ) + + transport.progress() + assert [request.test_count for request in comm.requests] == [1, 1, 1, 0, 0, 0] + transport.progress() + assert [request.test_count for request in comm.requests] == [1, 1, 1, 1, 1, 1] + transport.progress() + assert [request.test_count for request in comm.requests] == [2, 2, 2, 1, 1, 1] + + +def test_mpi_transport_rotates_receive_sources_under_a_busy_sender() -> None: + comm = _FakeMpiComm() + transport = _make_unit_mpi_transport(comm) + for source in (0, 0, 0, 1, 2): + comm.receive_queues[source].append( + _Packet( + _MessageKind.VOTE, + ConsensusPhase.TERMINAL, + 800 + source, + 0, + ConsensusOutcome.COMPLETED, + source, + ) + ) + + observed_sources = [transport.receive(1)[0].source for _ in range(3)] + assert observed_sources == [0, 1, 2] + + +def test_shutdown_drains_all_ranks_and_rejects_future_publication() -> None: + _, transports, coordinators = _make_group() + + coordinators[0].publish_terminal(301, ConsensusOutcome.COMPLETED) + with ThreadPoolExecutor(max_workers=len(coordinators)) as executor: + futures = [executor.submit(instance.shutdown, 2.0) for instance in coordinators] + for future in futures: + future.result(timeout=3.0) + + assert all(transport.close_timeouts for transport in transports) + with pytest.raises(RuntimeError, match="after shutdown starts"): + coordinators[0].publish_ready(302) + + +def test_shutdown_retry_does_not_duplicate_close_request() -> None: + network, _, coordinators = _make_group() + follower = coordinators[0] + coordinator = coordinators[-1] + + with pytest.raises(RuntimeError, match="shutdown acknowledgement"): + follower.shutdown(0.001) + assert network.queued(coordinator.rank) == 1 + + with ThreadPoolExecutor(max_workers=len(coordinators)) as executor: + futures = [executor.submit(instance.shutdown, 2.0) for instance in coordinators] + for future in futures: + future.result(timeout=3.0) + + assert network.queued(coordinator.rank) == 0 diff --git a/tests/unittest/disaggregated/test_async_consensus_mpi.py b/tests/unittest/disaggregated/test_async_consensus_mpi.py new file mode 100644 index 000000000000..fdd652f06889 --- /dev/null +++ b/tests/unittest/disaggregated/test_async_consensus_mpi.py @@ -0,0 +1,65 @@ +# 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. + +"""Subprocess test for the real four-rank asynchronous consensus transport.""" + +from __future__ import annotations + +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest + +_PROCESS_TIMEOUT_S = 90 + + +@pytest.mark.timeout(_PROCESS_TIMEOUT_S) +def test_async_consensus_real_mpi_four_rank_protocol() -> None: + pytest.importorskip("mpi4py") + mpirun = shutil.which("mpirun") + if mpirun is None: + pytest.skip("mpirun not found on PATH") + + worker = Path(__file__).with_name("async_consensus_mpi_worker.py") + env = os.environ.copy() + env.update( + { + "OMPI_ALLOW_RUN_AS_ROOT": "1", + "OMPI_ALLOW_RUN_AS_ROOT_CONFIRM": "1", + "PYTHONUNBUFFERED": "1", + } + ) + completed = subprocess.run( + [ + mpirun, + "--allow-run-as-root", + "--oversubscribe", + "-np", + "4", + sys.executable, + str(worker), + ], + check=False, + capture_output=True, + env=env, + text=True, + timeout=_PROCESS_TIMEOUT_S - 5, + ) + output = completed.stdout + completed.stderr + assert completed.returncode == 0, output + assert "ASYNC_CONSENSUS_MPI_OK" in output, output diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index be94d873ed61..c2b85fdb5fed 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -152,6 +152,14 @@ def test_no_tail_returns_none(self): # non-bounced result: only [msg_type, prefix] -> no tail. assert btr.decode_result_tail([b"KV_AGENT_RESULT", b"prefix"]) == (None, None, None) + def test_v2_roundtrip_uses_explicit_tail_index(self): + wm = self._wm([100], [16], 0xABCD) + msg = [b"KV_AGENT_RESULT", b"prefix", b"sender-endpoint"] + btr.encode_result_tail(wm) + dst, sizes, src_base = btr.decode_result_tail(msg, tail_index=3) + assert np.array_equal(dst, wm.dst_ptrs) + assert np.array_equal(sizes, wm.sizes) + assert src_base == 0xABCD + def test_encode_tail_handles_unset_base(self): wm = self._wm([1, 2], [8, 8], None) tail = btr.encode_result_tail(wm) @@ -187,6 +195,30 @@ def test_make_kv_result_msg_uses_binary_frame(result_name): assert tfr._AGENT_RESULT_BY_CODE[code] is result +def test_make_kv_result_msg_v2_preserves_identity_size_and_tail(): + """The v2 prefix adds identity without dropping timing or bounce framing.""" + tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") + tail = [b"dst", b"sizes", b"base"] + msg = tfr._make_kv_result_msg( + 3, + 12345, + 7, + True, + tfr.AgentResult.SUCCESS, + transfer_size=8192, + tail=tail, + request_epoch=99, + sender_endpoint="tcp://sender", + ) + + assert msg[0] == tfr.MessageType.KV_AGENT_RESULT + assert msg[2] == b"tcp://sender" + assert msg[3:] == tail + rank, rid, epoch, slice_id, last, code, size = tfr._KV_RESULT_PREFIX_V2.unpack(msg[1]) + assert (rank, rid, epoch, slice_id, last, size) == (3, 12345, 99, 7, True, 8192) + assert tfr._AGENT_RESULT_BY_CODE[code] is tfr.AgentResult.SUCCESS + + # --------------------------------------------------------------------------- # # fan-in safety gate — equal total//num_writers split only for uniform TP-by-head # --------------------------------------------------------------------------- # @@ -232,6 +264,39 @@ def pt(mapper_kind): assert safe(ov(1, 1, ranks=(0, 1)), ri([24], pt(MapperKind.NHD))) is True +def test_receiver_derives_canonical_rank_bound_destination_plan(): + tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") + if not hasattr(tfr.Receiver, "_build_bounce_destination_plan"): + pytest.skip("minimal local transfer stub does not load Receiver implementation") + receiver = object.__new__(tfr.Receiver) + local_region = SimpleNamespace(memory=SimpleNamespace(ptrs=np.array([1, 2]))) + peer_region = SimpleNamespace(memory=SimpleNamespace(ptrs=np.array([3, 4]))) + mapped = SimpleNamespace( + src=SimpleNamespace( + memory=SimpleNamespace( + ptrs=np.array([0x2050, 0x2000], dtype=np.int64), + bytes_per_region=0x50, + ) + ) + ) + mapper = SimpleNamespace(map=lambda _local, _peer: mapped) + receiver._registrar = SimpleNamespace( + self_extractor=SimpleNamespace(extract=lambda *_args, **_kwargs: local_region), + peer_extractor=lambda *_args: SimpleNamespace( + extract=lambda *_extract_args, **_extract_kwargs: peer_region + ), + get_pool_mapping=lambda _peer: {(0, 0): (0, 0)}, + get_kv_map=lambda *_args: mapper, + ) + receiver_req = SimpleNamespace(block_ids_per_layer_groups=[np.array([5, 6], dtype=np.int64)]) + peer_ri = SimpleNamespace(instance_name="ctx", instance_rank=7) + + dst_ptrs, sizes = receiver._build_bounce_destination_plan(receiver_req, peer_ri) + + assert dst_ptrs.tolist() == [0x2000, 0x2050] + assert sizes.tolist() == [0x50, 0x50] + + # --------------------------------------------------------------------------- # # NoBounceTransport — disabled no-op fallback # --------------------------------------------------------------------------- # @@ -244,6 +309,7 @@ def test_noop_behaviour(self): assert nb.build_request(SimpleNamespace()) is None assert nb.writer_base(("r", 0), 1) is None assert nb.is_bounced(("r", 0)) is False + nb.set_completion_callback(("r", 0), lambda _ok: None) nb.record_result(("r", 0), 1) # no-op, must not raise nb.record_failure(("r", 0), 1) # no-op, must not raise nb.release_idle_reservation(("r", 0)) # no-op, must not raise @@ -301,7 +367,14 @@ def reg_descs(self): return [] -def _make_transport(monkeypatch, block_bytes_per_group, capacity=1 << 30, min_blocks=1): +def _make_transport( + monkeypatch, + block_bytes_per_group, + capacity=1 << 30, + min_blocks=1, + destination_pool_layouts=None, + valid_destination_ranges=None, +): monkeypatch.setattr(btr, "SlotAllocator", _FakeAlloc) monkeypatch.setattr(btr.VmmBounceTransport, "_new_stream", lambda self: 0) monkeypatch.setattr( @@ -316,6 +389,10 @@ def _make_transport(monkeypatch, block_bytes_per_group, capacity=1 << 30, min_bl capacity_bytes=capacity, phys_chunk_size=32 * _MIB, block_bytes_per_group=block_bytes_per_group, + destination_pool_layouts=destination_pool_layouts, + valid_destination_ranges=( + [(0, 1 << 62)] if valid_destination_ranges is None else valid_destination_ranges + ), min_blocks=min_blocks, ) @@ -329,6 +406,18 @@ def _recv_req(block_counts, rid=1, slice_id=0): ) +def _recv_req_with_ids(block_ids_per_group, rid=1, slice_id=0): + return SimpleNamespace( + block_ids_per_layer_groups=[ + np.asarray(block_ids, dtype=np.int64) for block_ids in block_ids_per_group + ], + unique_rid=rid, + slice_id=slice_id, + bounce_dst_base=None, + mamba_state_index=None, + ) + + @pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.transport import needs CUDA bindings") class TestFanInReserve: def test_reserve_stamps_base_and_per_writer(self, monkeypatch): @@ -386,13 +475,15 @@ def test_fanin_scatters_ordered_by_src_base(self, monkeypatch): req = _recv_req([2]) assert t.reserve(req, num_writers=2) is True rid_slice = (req.unique_rid, req.slice_id) + low_base = t.bind_writer(rid_slice, 3, 0) + high_base = t.bind_writer(rid_slice, 7, 1) # writer for the HIGHER src_base reports first; scatter must reorder by src_base. t.record_result( rid_slice, 7, np.array([20], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=200, + np.array([100], dtype=np.int64), + src_base=high_base, ) assert t._scatter_q.empty() # only 1 of 2 writers terminal -> no scatter assert not t._recv_alloc.released # region NOT freed while a writer is still pending @@ -400,36 +491,29 @@ def test_fanin_scatters_ordered_by_src_base(self, monkeypatch): rid_slice, 3, np.array([10], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=100, + np.array([100], dtype=np.int64), + src_base=low_base, ) ctx, descs = t._scatter_q.get_nowait() - # each tail carries its OWN src_base; sorted (100 before 200) so the scatter is deterministic. - assert [t[0] for t in descs] == [100, 200] # per-writer src_base preserved + # Each tail carries its own bound source; sorting makes scatter deterministic. + assert [item[0] for item in descs] == [low_base, high_base] assert [list(t[1]) for t in descs] == [[10], [20]] # dst_ptrs - assert [list(t[2]) for t in descs] == [[8], [8]] # sizes + assert [list(t[2]) for t in descs] == [[100], [100]] # sizes - def test_fanin_fallback_writer_leaves_survivor_at_its_own_base(self, monkeypatch): - # Regression: if one fan-in writer falls back to in-place (SUCCESS, empty tail) while a - # sibling bounces, the survivor must be scattered from ITS OWN src_base, not packed to 0. + def test_incomplete_bounced_result_is_rejected_without_releasing(self, monkeypatch): + # Once the receiver advertises a bounce address, a tail-less SUCCESS is malformed rather + # than an in-place fallback. It must not consume writer credit or free the shared region. t = _make_transport(monkeypatch, block_bytes_per_group=[100]) req = _recv_req([2]) assert t.reserve(req, num_writers=2) is True rid_slice = (req.unique_rid, req.slice_id) - t.record_result( - rid_slice, 7, None, None - ) # writer 0 fell back to in-place: SUCCESS, no tail - assert t._scatter_q.empty() # only 1 of 2 writers terminal - t.record_result( - rid_slice, - 3, - np.array([10], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=100, - ) # writer 1 bounced to base+100 - ctx, descs = t._scatter_q.get_nowait() - assert [t[0] for t in descs] == [100] # only the survivor, read from base+100 (NOT 0) - assert [list(t[1]) for t in descs] == [[10]] + t.bind_writer(rid_slice, 7, 0) + t.bind_writer(rid_slice, 3, 1) + with pytest.raises(RuntimeError, match="incomplete bounced-result scatter tail"): + t.record_result(rid_slice, 7, None, None) + assert t._scatter_q.empty() + assert t._recv_alloc.released == [] + assert t.is_bounced(rid_slice) is True def test_fanin_failed_then_success_releases_only_after_both(self, monkeypatch): # A FAILED writer must not free the shared region until every writer is terminal. @@ -437,11 +521,17 @@ def test_fanin_failed_then_success_releases_only_after_both(self, monkeypatch): req = _recv_req([2]) assert t.reserve(req, num_writers=2) is True rid_slice = (req.unique_rid, req.slice_id) + t.bind_writer(rid_slice, 7, 0) + src_base = t.bind_writer(rid_slice, 3, 1) t.record_failure(rid_slice, 7) # first writer fails assert not t._recv_alloc.released # region held while a sibling may still be in flight assert t.is_bounced(rid_slice) is True t.record_result( - rid_slice, 3, np.array([10], dtype=np.int64), np.array([8], dtype=np.int64), src_base=0 + rid_slice, + 3, + np.array([10], dtype=np.int64), + np.array([100], dtype=np.int64), + src_base=src_base, ) # all terminal, >=1 FAILED -> no scatter, release (both drained), region freed. assert t._scatter_q.empty() @@ -456,13 +546,14 @@ def test_on_done_fires_after_scatter_lands(self, monkeypatch): req = _recv_req([2]) assert t.reserve(req, num_writers=1) is True rid_slice = (req.unique_rid, req.slice_id) + src_base = t.bind_writer(rid_slice, 3, 0) calls = [] t.record_result( rid_slice, 3, np.array([10], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=0, + np.array([200], dtype=np.int64), + src_base=src_base, on_done=lambda ok: calls.append(ok), ) ctx, descs = t._scatter_q.get_nowait() @@ -473,33 +564,36 @@ def test_on_done_fires_after_scatter_lands(self, monkeypatch): assert calls == [True] assert t._recv_alloc.released - def test_empty_acc_fires_on_done_inline_and_releases(self, monkeypatch): - # Bounced SUCCESS that carried no scatter tail: nothing to copy, but the task must still - # complete -> on_done(True) inline + slot released, nothing queued. + def test_empty_tail_does_not_fire_callback_or_release(self, monkeypatch): + # A missing scatter plan is not proof that the advertised remote write landed correctly. t = _make_transport(monkeypatch, block_bytes_per_group=[100]) req = _recv_req([2]) assert t.reserve(req, num_writers=1) is True rid_slice = (req.unique_rid, req.slice_id) + t.bind_writer(rid_slice, 3, 0) calls = [] - t.record_result(rid_slice, 3, None, None, on_done=lambda ok: calls.append(ok)) - assert calls == [True] + with pytest.raises(RuntimeError, match="incomplete bounced-result scatter tail"): + t.record_result(rid_slice, 3, None, None, on_done=lambda ok: calls.append(ok)) + assert calls == [] assert t._scatter_q.empty() - assert t._recv_alloc.released # slot freed + assert t._recv_alloc.released == [] + assert t.is_bounced(rid_slice) is True - def test_missing_key_is_dropped(self, monkeypatch): - # A late/duplicate result for an already-settled (popped) rid_slice is dropped: the context's own - # settle already fired completion, so re-firing here would double-report. + def test_missing_key_is_rejected(self, monkeypatch): + # Session tombstones filter legitimate delayed duplicates before they reach the bounce layer; + # an unknown reservation here is a protocol error and cannot be silently accepted. t = _make_transport(monkeypatch, block_bytes_per_group=[100]) calls = [] - t.record_result( - ("missing", 0), - 3, - np.array([10], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=0, - on_done=lambda ok: calls.append(ok), - ) - assert calls == [] # no-op, no callback + with pytest.raises(RuntimeError, match="unknown reservation"): + t.record_result( + ("missing", 0), + 3, + np.array([10], dtype=np.int64), + np.array([100], dtype=np.int64), + src_base=0, + on_done=lambda ok: calls.append(ok), + ) + assert calls == [] def test_duplicate_writer_is_ignored(self, monkeypatch): # A duplicate SUCCESS from the same peer_rank must not double-count toward all-terminal. @@ -507,12 +601,314 @@ def test_duplicate_writer_is_ignored(self, monkeypatch): req = _recv_req([2]) assert t.reserve(req, num_writers=2) is True rid_slice = (req.unique_rid, req.slice_id) - arr = (np.array([10], dtype=np.int64), np.array([8], dtype=np.int64)) - t.record_result(rid_slice, 7, *arr, src_base=0) - t.record_result(rid_slice, 7, *arr, src_base=0) # duplicate of the SAME writer + src_base = t.bind_writer(rid_slice, 7, 0) + t.bind_writer(rid_slice, 3, 1) + arr = (np.array([10], dtype=np.int64), np.array([100], dtype=np.int64)) + t.record_result(rid_slice, 7, *arr, src_base=src_base) + t.record_result(rid_slice, 7, *arr, src_base=src_base) # duplicate of the SAME writer assert t._scatter_q.empty() # still only 1 distinct writer -> not all terminal assert not t._recv_alloc.released + @pytest.mark.parametrize( + "peer_rank,source_offset,error", + [ + (9, 0, "source identity mismatch"), + (7, 1, "source identity mismatch"), + ], + ) + def test_unbound_or_wrong_source_writer_is_rejected( + self, monkeypatch, peer_rank, source_offset, error + ): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + req = _recv_req([1]) + assert t.reserve(req, num_writers=1) + rid_slice = (req.unique_rid, req.slice_id) + src_base = t.bind_writer(rid_slice, 7, 0) + with pytest.raises(RuntimeError, match=error): + t.record_result( + rid_slice, + peer_rank, + np.array([10], dtype=np.int64), + np.array([100], dtype=np.int64), + src_base=src_base + source_offset, + ) + assert t._recv_alloc.released == [] + assert t.is_bounced(rid_slice) + + @pytest.mark.parametrize( + "dst_ptrs,sizes,error", + [ + ([0x2000], [99], "describes 99 bytes"), + ([0x2000, 0x2030], [60, 40], "overlap"), + ([0x1FF0], [100], "outside the receiver-owned KV destination plan"), + ([0x2000], [0], "must be positive"), + ], + ) + def test_untrusted_scatter_plan_is_rejected(self, monkeypatch, dst_ptrs, sizes, error): + t = _make_transport( + monkeypatch, + block_bytes_per_group=[100], + valid_destination_ranges=[(0x2000, 0x2100)], + ) + req = _recv_req([1]) + assert t.reserve(req, num_writers=1) + rid_slice = (req.unique_rid, req.slice_id) + src_base = t.bind_writer(rid_slice, 7, 0) + with pytest.raises(RuntimeError, match=error): + t.record_result( + rid_slice, + 7, + np.array(dst_ptrs, dtype=np.int64), + np.array(sizes, dtype=np.int64), + src_base=src_base, + ) + assert t._recv_alloc.released == [] + assert t.is_bounced(rid_slice) + + def test_scatter_plan_cannot_target_another_request_block(self, monkeypatch): + # Both blocks live in the registered pool, but this request owns only + # block 2. A sender-returned tail naming block 3 must fail closed. + t = _make_transport( + monkeypatch, + block_bytes_per_group=[100], + destination_pool_layouts=[[(0x2000, 100, 8)]], + valid_destination_ranges=[(0x2000, 0x2320)], + ) + req = _recv_req_with_ids([[2]]) + expected_ptr = 0x2000 + 2 * 100 + assert t.reserve( + req, + num_writers=1, + expected_destination_plans={ + 7: ( + np.array([expected_ptr], dtype=np.int64), + np.array([100], dtype=np.int64), + ) + }, + ) + rid_slice = (req.unique_rid, req.slice_id) + src_base = t.bind_writer(rid_slice, 7, 0) + + with pytest.raises(RuntimeError, match="outside the receiver-owned KV destination plan"): + t.record_result( + rid_slice, + 7, + np.array([0x2000 + 3 * 100], dtype=np.int64), + np.array([100], dtype=np.int64), + src_base=src_base, + ) + + assert t._recv_alloc.released == [] + assert t.is_bounced(rid_slice) + + @pytest.mark.parametrize( + "dst_ptrs,error", + [ + ([0x2000, 0x2000], "overlap or duplicate"), + ([0x2032, 0x2000], "not in canonical address order"), + ], + ) + def test_exact_plan_rejects_duplicate_or_reordered_tail(self, monkeypatch, dst_ptrs, error): + t = _make_transport( + monkeypatch, + block_bytes_per_group=[100], + destination_pool_layouts=[[(0x2000, 100, 8)]], + ) + req = _recv_req_with_ids([[0]]) + expected = { + 7: ( + np.array([0x2000, 0x2032], dtype=np.int64), + np.array([50, 50], dtype=np.int64), + ) + } + assert t.reserve(req, expected_destination_plans=expected) + rid_slice = (req.unique_rid, req.slice_id) + src_base = t.bind_writer(rid_slice, 7, 0) + + with pytest.raises(RuntimeError, match=error): + t.record_result( + rid_slice, + 7, + np.array(dst_ptrs, dtype=np.int64), + np.array([50, 50], dtype=np.int64), + src_base=src_base, + ) + + assert t._scatter_q.empty() + assert t._recv_alloc.released == [] + assert t.is_bounced(rid_slice) + + @pytest.mark.parametrize("wrong_ptr", [0x2001, 0x2064]) + def test_exact_plan_rejects_rank_swapped_or_in_bounds_wrong_offset( + self, monkeypatch, wrong_ptr + ): + # Both writer plans are valid, in-request ranges and collectively cover + # the slot. A writer still cannot claim its sibling's range or shift its + # own range within the slot: rank identity binds the exact sequence. + t = _make_transport( + monkeypatch, + block_bytes_per_group=[200], + destination_pool_layouts=[[(0x2000, 200, 8)]], + ) + req = _recv_req_with_ids([[0]]) + expected = { + 7: ( + np.array([0x2000], dtype=np.int64), + np.array([100], dtype=np.int64), + ), + 3: ( + np.array([0x2064], dtype=np.int64), + np.array([100], dtype=np.int64), + ), + } + assert t.reserve(req, num_writers=2, expected_destination_plans=expected) + rid_slice = (req.unique_rid, req.slice_id) + src_base = t.bind_writer(rid_slice, 7, 0) + t.bind_writer(rid_slice, 3, 1) + + with pytest.raises(RuntimeError, match="exact receiver-derived destination plan"): + t.record_result( + rid_slice, + 7, + np.array([wrong_ptr], dtype=np.int64), + np.array([100], dtype=np.int64), + src_base=src_base, + ) + + assert t._scatter_q.empty() + assert t._recv_alloc.released == [] + assert t.is_bounced(rid_slice) + + def test_exact_rank_bound_plans_accept_correct_fanin(self, monkeypatch): + t = _make_transport( + monkeypatch, + block_bytes_per_group=[200], + destination_pool_layouts=[[(0x2000, 200, 8)]], + ) + req = _recv_req_with_ids([[0]]) + expected = { + 7: ( + np.array([0x2000], dtype=np.int64), + np.array([100], dtype=np.int64), + ), + 3: ( + np.array([0x2064], dtype=np.int64), + np.array([100], dtype=np.int64), + ), + } + assert t.reserve(req, num_writers=2, expected_destination_plans=expected) + rid_slice = (req.unique_rid, req.slice_id) + first_base = t.bind_writer(rid_slice, 7, 0) + second_base = t.bind_writer(rid_slice, 3, 1) + + t.record_result( + rid_slice, + 3, + np.array([0x2064], dtype=np.int64), + np.array([100], dtype=np.int64), + src_base=second_base, + ) + t.record_result( + rid_slice, + 7, + np.array([0x2000], dtype=np.int64), + np.array([100], dtype=np.int64), + src_base=first_base, + ) + + _ctx, descs = t._scatter_q.get_nowait() + assert [src for src, _dst, _sizes in descs] == [first_base, second_base] + + def test_request_destination_plan_rejects_out_of_range_block_id(self, monkeypatch): + t = _make_transport( + monkeypatch, + block_bytes_per_group=[100], + destination_pool_layouts=[[(0x2000, 100, 8)]], + ) + req = _recv_req_with_ids([[8]]) + + assert t.reserve(req, num_writers=1) is False + assert req.bounce_dst_base is None + assert t._recv_alloc.released == [0] + + def test_mamba_request_does_not_advertise_an_incomplete_bounce_plan(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + req = _recv_req_with_ids([[0]]) + req.mamba_state_index = 4 + + assert t.reserve(req, num_writers=1) is False + assert req.bounce_dst_base is None + + def test_build_request_releases_send_slot_when_request_creation_fails(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + monkeypatch.setattr(t, "_gather_blocking", lambda *args, **kwargs: None) + + def fail_make_write(*args, **kwargs): + raise RuntimeError("descriptor creation failed") + + monkeypatch.setattr(t, "_make_write", fail_make_write) + write_meta = SimpleNamespace( + src_ptrs=np.array([0x1000], dtype=np.int64), + dst_ptrs=np.array([0x2000], dtype=np.int64), + sizes=np.array([100], dtype=np.int64), + ) + with pytest.raises(RuntimeError, match="descriptor creation failed"): + t.build_request(write_meta) + assert t._send_alloc.released == [0] + + def test_build_request_canonicalizes_coupled_fragment_triplets(self, monkeypatch): + t = _make_transport(monkeypatch, block_bytes_per_group=[100]) + gathered = [] + monkeypatch.setattr( + t, + "_gather_blocking", + lambda _addr, meta, _total: gathered.append( + (meta.src_ptrs.copy(), meta.dst_ptrs.copy(), meta.sizes.copy()) + ), + ) + monkeypatch.setattr(t, "_make_write", lambda *_args: "request") + write_meta = SimpleNamespace( + src_ptrs=np.array([0x3000, 0x1000, 0x2000], dtype=np.int64), + dst_ptrs=np.array([0x2300, 0x2100, 0x2200], dtype=np.int64), + sizes=np.array([30, 10, 20], dtype=np.int64), + ) + + request, _slot_id = t.build_request(write_meta) + + assert request == "request" + assert list(write_meta.src_ptrs) == [0x1000, 0x2000, 0x3000] + assert list(write_meta.dst_ptrs) == [0x2100, 0x2200, 0x2300] + assert list(write_meta.sizes) == [10, 20, 30] + assert [list(values) for values in gathered[0]] == [ + [0x1000, 0x2000, 0x3000], + [0x2100, 0x2200, 0x2300], + [10, 20, 30], + ] + + def test_close_retains_memory_if_scatter_thread_does_not_exit(self): + t = object.__new__(btr.VmmBounceTransport) + stopped = [] + joined = [] + deregistered = [] + closed = [] + t._stop = SimpleNamespace(set=lambda: stopped.append(True)) + t._scatter_q = queue.Queue() + t._scatter_thread = SimpleNamespace( + is_alive=lambda: True, join=lambda timeout: joined.append(timeout) + ) + t._reg_descs = ["send", "recv"] + t._agent = SimpleNamespace(deregister_memory=lambda desc: deregistered.append(desc)) + t._send_alloc = SimpleNamespace(close=lambda: closed.append("send")) + t._recv_alloc = SimpleNamespace(close=lambda: closed.append("recv")) + + with pytest.raises(RuntimeError, match="did not exit"): + t.close() + + assert stopped == [True] + assert joined == [btr._CLOSE_JOIN_S] + assert deregistered == [] + assert closed == [] + def test_scatter_write_result_non_bounce_fires_on_done(self): # Non-bounced path completes inline (the in-place WRITE already landed the KV). calls = [] @@ -533,19 +929,26 @@ def test_release_idle_reservation_frees_slot_and_is_idempotent(self, monkeypatch assert t._recv_alloc.released # slot freed t.release_idle_reservation(rid_slice) # already gone -> no-op, must not raise - def test_orphan_reservation_quarantines_and_is_idempotent(self, monkeypatch): - # Giving up on an in-flight reservation must quarantine the region (a write may still land), - # not release or leak it; a second give-up is a no-op. + def test_orphan_reservation_waits_for_explicit_drain_proof(self, monkeypatch): + # A fixed quarantine timeout cannot prove a remote RMA is done. Keep + # the slot live until the sender ACK supplies exact drain proof. t = _make_transport(monkeypatch, block_bytes_per_group=[100]) req = _recv_req([2]) assert t.reserve(req, num_writers=1) is True rid_slice = (req.unique_rid, req.slice_id) t.orphan_reservation(rid_slice) - assert t._recv_alloc.quarantined == [0] # quarantined, not released + assert t._recv_alloc.quarantined == [] assert t._recv_alloc.released == [] - assert t.is_bounced(rid_slice) is False # settled and removed from the live map - t.orphan_reservation(rid_slice) # already gone -> no-op, must not raise - assert t._recv_alloc.quarantined == [0] + assert t.is_bounced(rid_slice) is True + + t._recv_alloc.reclaim_expired() + assert t.is_bounced(rid_slice) is True + assert t._recv_alloc.released == [] + + t.confirm_drained(rid_slice) + assert t.is_bounced(rid_slice) is False + assert t._recv_alloc.released == [0] + t.confirm_drained(rid_slice) # already gone -> idempotent # --------------------------------------------------------------------------- # @@ -597,7 +1000,7 @@ def _ctx(self, num_writers, per_writer_bytes=100, base_addr=0x1000): ) def _dst(self, v=10): - return dict(dst_ptrs=np.array([v], dtype=np.int64), sizes=np.array([8], dtype=np.int64)) + return dict(dst_ptrs=np.array([v], dtype=np.int64), sizes=np.array([100], dtype=np.int64)) def test_writer_base_layout(self): c = self._ctx(3, per_writer_bytes=0x64, base_addr=0x1000) @@ -605,8 +1008,9 @@ def test_writer_base_layout(self): def test_single_writer_success_scatters_then_releases(self): c = self._ctx(1) + c.bind_writer(3, c.writer_base(0)) assert not c.ready_to_scatter() and not c.ready_to_settle() - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) + c.record_writer_result(3, succeeded=True, src_base=c.writer_base(0), **self._dst()) assert c.ready_to_scatter() c.begin_scatter() assert not c.ready_to_settle() # scatter not landed yet @@ -619,17 +1023,21 @@ def test_single_writer_success_scatters_then_releases(self): def test_fanin_holds_until_all_terminal(self): c = self._ctx(2) - c.record_writer_result(7, succeeded=True, src_base=0, **self._dst()) + c.bind_writer(7, c.writer_base(0)) + c.bind_writer(3, c.writer_base(1)) + c.record_writer_result(7, succeeded=True, src_base=c.writer_base(0), **self._dst()) assert not c.ready_to_scatter() # 1/2 writers assert not c.ready_to_settle() # drain-before-release - c.record_writer_result(3, succeeded=True, src_base=100, **self._dst()) + c.record_writer_result(3, succeeded=True, src_base=c.writer_base(1), **self._dst()) assert c.ready_to_scatter() # all success -> scatter def test_fanin_failed_then_success_releases(self): c = self._ctx(2) + c.bind_writer(7, c.writer_base(0)) + c.bind_writer(3, c.writer_base(1)) c.record_writer_result(7, succeeded=False) assert not c.ready_to_settle() # a sibling is still pending -> hold - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) + c.record_writer_result(3, succeeded=True, src_base=c.writer_base(1), **self._dst()) assert not c.ready_to_scatter() # >=1 FAILED -> skip scatter assert c.ready_to_settle() ret = c.settle() @@ -637,39 +1045,101 @@ def test_fanin_failed_then_success_releases(self): assert ret.success is False assert c.state is bcore.TransferState.FAILED - def test_orphan_quarantines(self): + def test_orphan_waits_for_drain_proof(self): c = self._ctx(2) - c.record_writer_result(7, succeeded=True, src_base=0, **self._dst()) + c.bind_writer(7, c.writer_base(0)) + c.bind_writer(3, c.writer_base(1)) + c.record_writer_result(7, succeeded=True, src_base=c.writer_base(0), **self._dst()) c.mark_orphaned() # the other writer is in-doubt + assert not c.ready_to_settle() + c.confirm_drained() assert c.ready_to_settle() ret = c.settle() - assert ret.disposition is bcore.Disposition.QUARANTINE and ret.success is False - assert c.state is bcore.TransferState.QUARANTINED + assert ret.disposition is bcore.Disposition.RELEASE and ret.success is False + assert c.state is bcore.TransferState.CANCELLED_DRAINED - def test_empty_tail_success_releases_without_scatter(self): + def test_orphan_drain_proof_fires_unconditional_settlement_callback(self): + calls = [] c = self._ctx(1) - c.record_writer_result(3, succeeded=True) # no dst tail - assert not c.ready_to_scatter() - assert c.ready_to_settle() - assert c.settle().disposition is bcore.Disposition.RELEASE + c.bind_writer(3, c.writer_base(0)) + c.set_completion_callback(lambda ok: calls.append(("settled", ok))) + c.mark_orphaned() + c.confirm_drained() + + settlement = c.settle() + + assert settlement is not None + assert calls == [] + settlement.on_done(settlement.success) + assert calls == [("settled", False)] + + def test_success_fires_result_then_unconditional_settlement_callback(self): + calls = [] + c = self._ctx(1) + c.bind_writer(3, c.writer_base(0)) + c.on_done = lambda ok: calls.append(("result", ok)) + c.set_completion_callback(lambda ok: calls.append(("settled", ok))) + c.record_writer_result(3, succeeded=True, src_base=c.writer_base(0), **self._dst()) + c.begin_scatter() + c.finish_scatter(True) + + settlement = c.settle() + + assert settlement is not None + settlement.on_done(settlement.success) + assert calls == [("result", True), ("settled", True)] + + def test_settlement_callback_runs_when_result_callback_raises(self): + calls = [] + c = self._ctx(1) + c.bind_writer(3, c.writer_base(0)) + + def fail_result(_ok): + calls.append("result") + raise RuntimeError("boom") + + c.on_done = fail_result + c.set_completion_callback(lambda _ok: calls.append("settled")) + c.record_writer_result(3, succeeded=True, src_base=c.writer_base(0), **self._dst()) + c.begin_scatter() + c.finish_scatter(True) + + settlement = c.settle() + + assert settlement is not None + with pytest.raises(RuntimeError, match="boom"): + settlement.on_done(settlement.success) + assert calls == ["result", "settled"] + + def test_empty_tail_success_is_rejected(self): + c = self._ctx(1) + c.bind_writer(3, c.writer_base(0)) + with pytest.raises(RuntimeError, match="complete scatter tail"): + c.record_writer_result(3, succeeded=True) + assert not c.ready_to_settle() - def test_writers_locked_after_scatter_drops_late_writer(self): + def test_writers_locked_after_scatter_rejects_unbound_late_writer(self): c = self._ctx(1) - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) + c.bind_writer(3, c.writer_base(0)) + c.record_writer_result(3, succeeded=True, src_base=c.writer_base(0), **self._dst()) c.begin_scatter() # SCATTERING -> frozen - c.record_writer_result(9, succeeded=False) # a late / reordered report - assert 9 not in c._writer_ok # dropped, cannot re-arm the state + with pytest.raises(RuntimeError, match="unexpected bounce writer"): + c.record_writer_result(9, succeeded=False) + assert 9 not in c._writer_ok def test_duplicate_writer_dedup(self): c = self._ctx(2) - c.record_writer_result(7, succeeded=True, src_base=0, **self._dst()) + c.bind_writer(7, c.writer_base(0)) + c.bind_writer(3, c.writer_base(1)) + c.record_writer_result(7, succeeded=True, src_base=c.writer_base(0), **self._dst()) c.record_writer_result(7, succeeded=False) # same rank again -> ignored assert c._writer_ok[7] is True assert not c.ready_to_settle() # still only 1 distinct writer of 2 def test_scatter_failure_releases_as_failed(self): c = self._ctx(1) - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) + c.bind_writer(3, c.writer_base(0)) + c.record_writer_result(3, succeeded=True, src_base=c.writer_base(0), **self._dst()) c.begin_scatter() c.finish_scatter(False) # scatter kernel failed ret = c.settle() @@ -679,7 +1149,8 @@ def test_orphan_after_scatter_is_ignored(self): # once SCATTERING, all writers already reported SUCCESS -> nothing is in doubt, so a late # orphan (e.g. a racing cancel) must NOT downgrade a clean transfer to quarantine. c = self._ctx(1) - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) + c.bind_writer(3, c.writer_base(0)) + c.record_writer_result(3, succeeded=True, src_base=c.writer_base(0), **self._dst()) c.begin_scatter() c.mark_orphaned() # no-op after SCATTERING c.finish_scatter(True) diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index c2cfc1382e37..a3f77b0cc2bb 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -47,6 +47,7 @@ from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings import LayerType as LayerTypeCpp from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp +from tensorrt_llm.disaggregated_params import DisaggScheduleStyle from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.logger import logger @@ -637,6 +638,9 @@ def add_and_verify_request( valid_gen_transfer_workers = gen_transfer_workers unique_rid = uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF + schedule_style = ( + DisaggScheduleStyle.CONTEXT_FIRST if send_first else DisaggScheduleStyle.GENERATION_FIRST + ) ctx_request = LlmRequest( request_id=ctx_request_id, max_new_tokens=1, @@ -647,7 +651,10 @@ def add_and_verify_request( is_streaming=False, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, ) - ctx_request.py_disaggregated_params = DisaggregatedParams(disagg_request_id=unique_rid) + ctx_request.py_disaggregated_params = DisaggregatedParams( + disagg_request_id=unique_rid, + schedule_style=schedule_style, + ) ctx_request.add_new_token(8 + ctx_request_id, 0) ctx_request.py_draft_tokens = [ @@ -687,6 +694,7 @@ def add_and_verify_request( ctx_dp_rank=ctx_dp_rank, ctx_info_endpoint=ctx_info_endpoint, disagg_request_id=unique_rid, + schedule_style=schedule_style, ) # Add sequence to gen KV cache managers gen_kv_caches = [] diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index 04defdd1cd37..7cc402718cbf 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -138,6 +138,129 @@ def _make_completion_response( ) +@pytest.mark.asyncio +async def test_generation_first_refreshes_final_endpoint_and_strips_generation(): + service = _make_service("generation_first") + request = CompletionRequest(model="model", prompt=[1, 2, 3]) + service._coordinator.get_disagg_request_id = AsyncMock(return_value=42) + service._ctx_router.get_next_server = AsyncMock( + return_value=( + "ctx:8000", + { + "server_info": { + "disaggregated_params": { + "ctx_info_endpoint": ["tcp://stale:1000"], + "ctx_endpoint_generation": "stale-generation", + } + } + }, + ) + ) + service._ctx_router.get_runtime_server_info = AsyncMock( + return_value={ + "disaggregated_params": { + "ctx_dp_rank": 0, + "ctx_info_endpoint": ["tcp://final:2000"], + "ctx_endpoint_generation": "final-generation", + } + } + ) + captured_gen_request = None + + async def send_gen(gen_request, **_kwargs): + nonlocal captured_gen_request + captured_gen_request = gen_request + return _make_completion_response("done", "length", context_only=False) + + service._gen_client = SimpleNamespace(send_request=AsyncMock(side_effect=send_gen)) + service._ctx_client = SimpleNamespace(send_request=AsyncMock(return_value=None)) + + await service._send_disagg_request_gen_first(request) + + service._ctx_router.get_runtime_server_info.assert_awaited_once_with( + "ctx:8000", require_generation=True + ) + assert captured_gen_request.disaggregated_params.ctx_info_endpoint == "tcp://final:2000" + assert not hasattr(captured_gen_request.disaggregated_params, "ctx_endpoint_generation") + + +@pytest.mark.asyncio +async def test_generation_first_flags_off_uses_valid_cached_endpoint(): + service = _make_service("generation_first") + request = CompletionRequest(model="model", prompt=[1, 2, 3]) + service._coordinator.get_disagg_request_id = AsyncMock(return_value=45) + service._ctx_router.get_next_server = AsyncMock( + return_value=( + "ctx:8000", + { + "server_info": { + "disaggregated_params": { + "ctx_info_endpoint": ["tcp://legacy:1000"], + } + } + }, + ) + ) + service._ctx_router.get_runtime_server_info = AsyncMock() + service._ctx_client = SimpleNamespace(send_request=AsyncMock(return_value=None)) + service._gen_client = SimpleNamespace( + send_request=AsyncMock( + return_value=_make_completion_response("done", "length", context_only=False) + ) + ) + + await service._send_disagg_request_gen_first(request) + + service._ctx_router.get_runtime_server_info.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_generation_first_refresh_failure_releases_ctx_reservation(): + service = _make_service("generation_first") + request = CompletionRequest(model="model", prompt=[1, 2, 3]) + service._coordinator.get_disagg_request_id = AsyncMock(return_value=43) + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:8000", {"server_info": {}})) + service._ctx_router.get_runtime_server_info = AsyncMock(side_effect=RuntimeError("not final")) + service._ctx_router.finish_request = AsyncMock() + + with pytest.raises(RuntimeError, match="not final"): + await service._send_disagg_request_gen_first(request) + + service._ctx_router.finish_request.assert_awaited_once_with(request, success=False, req_id=43) + + +@pytest.mark.asyncio +async def test_generation_first_cancelled_refresh_still_releases_reservation(): + service = _make_service("generation_first") + request = CompletionRequest(model="model", prompt=[1, 2, 3]) + service._coordinator.get_disagg_request_id = AsyncMock(return_value=44) + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:8000", {"server_info": {}})) + refresh_started = asyncio.Event() + + async def refresh(*_args, **_kwargs): + refresh_started.set() + await asyncio.Event().wait() + + service._ctx_router.get_runtime_server_info = AsyncMock(side_effect=refresh) + cleanup_finished = asyncio.Event() + + async def finish(*_args, **_kwargs): + await asyncio.sleep(0) + cleanup_finished.set() + + service._ctx_router.finish_request = AsyncMock(side_effect=finish) + + service_task = asyncio.create_task(service._send_disagg_request_gen_first(request)) + await refresh_started.wait() + service_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await service_task + + await asyncio.wait_for(cleanup_finished.wait(), timeout=1.0) + service._ctx_router.finish_request.assert_awaited_once_with(request, success=False, req_id=44) + + def _make_chat_response( finish_reason: str, disagg_request_id: int = 42, diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index ddaf12a0c8e2..e84397d052ed 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -34,6 +34,90 @@ # yapf: enable +@pytest.mark.asyncio +async def test_runtime_server_info_refreshes_generation_and_coalesces(): + server = "ctx:8000" + router = RoundRobinRouter(server_role="context", servers=[server]) + first = asyncio.Event() + release = asyncio.Event() + responses = [ + { + "disaggregated_params": { + "ctx_info_endpoint": ["tcp://ctx:1000"], + "ctx_endpoint_generation": "generation-1", + } + }, + { + "disaggregated_params": { + "ctx_info_endpoint": ["tcp://ctx:2000"], + "ctx_endpoint_generation": "generation-2", + } + }, + ] + + async def fetch(_server, _timeout): + first.set() + await release.wait() + return responses.pop(0) + + router._fetch_server_info = mock.AsyncMock(side_effect=fetch) + refresh_1 = asyncio.create_task( + router.get_runtime_server_info(server, require_generation=True)) + await first.wait() + refresh_2 = asyncio.create_task( + router.get_runtime_server_info(server, require_generation=True)) + release.set() + + result_1, result_2 = await asyncio.gather(refresh_1, refresh_2) + assert result_1 == result_2 + assert router._fetch_server_info.await_count == 1 + + restarted = await router.get_runtime_server_info(server, + require_generation=True) + assert restarted["disaggregated_params"][ + "ctx_endpoint_generation"] == "generation-2" + assert router._fetch_server_info.await_count == 2 + + +@pytest.mark.asyncio +async def test_runtime_server_info_rejects_missing_final_endpoint(): + server = "ctx:8000" + router = RoundRobinRouter(server_role="context", servers=[server]) + router._fetch_server_info = mock.AsyncMock( + return_value={"disaggregated_params": {}}) + + with pytest.raises(RuntimeError, match="final-runtime server info"): + await router.get_runtime_server_info(server, require_generation=True) + + assert router._fetch_server_info.await_count == 3 + assert all(call.args[1] == 5.0 + for call in router._fetch_server_info.await_args_list) + + +@pytest.mark.asyncio +async def test_delegating_router_refresh_uses_local_worker_session(): + # Stateful delegating clients intentionally do not mirror the + # coordinator's dynamic server list locally. + local = RoundRobinRouter(server_role="context", servers=[]) + server_info = { + "disaggregated_params": { + "ctx_info_endpoint": ["tcp://ctx:1000"], + "ctx_endpoint_generation": "generation-1", + } + } + local._fetch_server_info = mock.AsyncMock(return_value=server_info) + router = CoordinatorDelegatingRouter("unix:/tmp/coordinator.sock", local, + "context") + + result = await router.get_runtime_server_info("ctx:8000", + require_generation=True) + + assert result == server_info + local._fetch_server_info.assert_awaited_once_with("ctx:8000", 5.0) + assert local._server_info == {} + assert local._runtime_server_info_logged_generations == {} + + def test_native_block_key_hasher_matches_python_v1(): """Native C++ BlockKeyHasher must be bit-exact with hash_v1_block_key. diff --git a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py index 68b38bf896ae..b67e44517ebc 100644 --- a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py +++ b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py @@ -16,30 +16,78 @@ from __future__ import annotations +import threading +from collections import OrderedDict, defaultdict from dataclasses import dataclass -from typing import Optional +from types import SimpleNamespace +from typing import Callable, Optional from unittest.mock import Mock import pytest +from tensorrt_llm._torch.disaggregation import transceiver as transceiver_module +from tensorrt_llm._torch.disaggregation.async_consensus import ( + ConsensusEvent, + ConsensusEventKind, + ConsensusOutcome, +) from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus, WaitResult from tensorrt_llm._torch.disaggregation.native.transfer import TaskStatus, TxSession from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 from tensorrt_llm.bindings import LlmRequestState +@pytest.fixture(autouse=True) +def _clear_async_consensus_env(monkeypatch) -> None: + monkeypatch.delenv(transceiver_module._ASYNC_TERMINAL_ENV, raising=False) + monkeypatch.delenv(transceiver_module._ASYNC_PEER_READY_ENV, raising=False) + monkeypatch.delenv(transceiver_module._CONTEXT_ACTIVATION_DIGEST_ENV, raising=False) + + @dataclass class _FakeRequest: state: Optional[LlmRequestState] = None + request_id: int = 0 + py_disaggregated_params: Optional[object] = None + kv_cache_transfer_start: Optional[int] = None + + def set_kv_cache_transfer_start(self, value: int) -> None: + self.kv_cache_transfer_start = value class _FakeTransferWorker: def __init__(self) -> None: self.sweep_count = 0 + self.ready_request_ids: set[int] = set() + self.pinned_request_ids: set[int] = set() + self.cancelled_peer_request_ids: set[int] = set() + self.cancel_peer_result = True + self.tx_session = None + self.rx_session = None def sweep_stale_req_infos(self) -> None: self.sweep_count += 1 + def has_all_peer_req_infos_for_send(self, rid: int) -> bool: + return rid in self.ready_request_ids + + def pin_peer_req_infos_for_send(self, rid: int) -> None: + self.pinned_request_ids.add(rid) + + def cancel_peer_req_infos_for_send(self, rid: int) -> bool: + if self.cancel_peer_result: + self.pinned_request_ids.discard(rid) + self.cancelled_peer_request_ids.add(rid) + return self.cancel_peer_result + + def create_tx_session(self, _req): + assert self.tx_session is not None + return self.tx_session + + def create_rx_session(self, _req): + assert self.rx_session is not None + return self.rx_session + class _FakeSession: def __init__( @@ -50,14 +98,20 @@ def __init__( status: SessionStatus = SessionStatus.READY, is_completed: bool = False, has_failed: bool = False, + has_transferring_tasks: bool = False, + seal_quiescent: Optional[bool] = None, ) -> None: self._rid = rid self._wait_result = wait_result self._status = status self._is_completed = is_completed self._has_failed = has_failed + self._has_transferring_tasks = has_transferring_tasks + self._seal_quiescent = seal_quiescent self.blocking_calls: list[bool] = [] self.closed = False + self.sealed = False + self.cancel_calls = 0 @property def disagg_request_id(self) -> int: @@ -77,10 +131,65 @@ def is_completed(self) -> bool: def has_failed(self) -> bool: return self._has_failed + def has_transferring_tasks(self) -> bool: + return self._has_transferring_tasks + + def seal_and_check_quiescent(self) -> bool: + self.sealed = True + if self._seal_quiescent is not None: + return self._seal_quiescent + return not self._has_transferring_tasks + + def cancel(self) -> None: + self.cancel_calls += 1 + self._status = SessionStatus.CANCELLED + def close(self) -> None: self.closed = True +class _FakeAsyncCoordinator: + def __init__(self, *, rank: int = 1, scheduling_rank: int = 0) -> None: + self.rank = rank + self.scheduling_rank = scheduling_rank + self.events: list[ConsensusEvent] = [] + self.terminal_votes: list[tuple[int, ConsensusOutcome, int]] = [] + self.ready_votes: list[tuple[int, int]] = [] + self.ready_acks: list[tuple[int, int]] = [] + self.ready_activation_acks: list[tuple[int, int]] = [] + self.ready_abort_acks: list[tuple[int, int]] = [] + self.ready_withdrawals: list[tuple[int, int]] = [] + self.withdraw_result = True + self.poll_count = 0 + self.poll_hook: Optional[Callable[["_FakeAsyncCoordinator"], None]] = None + + def poll(self) -> list[ConsensusEvent]: + self.poll_count += 1 + if self.poll_hook is not None: + self.poll_hook(self) + events, self.events = self.events, [] + return events + + def publish_terminal(self, rid: int, outcome: ConsensusOutcome, epoch: int) -> None: + self.terminal_votes.append((rid, outcome, epoch)) + + def publish_ready(self, rid: int, epoch: int) -> None: + self.ready_votes.append((rid, epoch)) + + def acknowledge_ready(self, rid: int, epoch: int) -> None: + self.ready_acks.append((rid, epoch)) + + def acknowledge_ready_activation(self, rid: int, epoch: int) -> None: + self.ready_activation_acks.append((rid, epoch)) + + def acknowledge_ready_abort(self, rid: int, epoch: int) -> None: + self.ready_abort_acks.append((rid, epoch)) + + def withdraw_ready(self, rid: int, epoch: int) -> bool: + self.ready_withdrawals.append((rid, epoch)) + return self.withdraw_result + + class _FakeTask: def __init__(self, status: TaskStatus, wait_result: bool = True) -> None: self.status = status @@ -99,7 +208,25 @@ def _make_transceiver( transceiver = object.__new__(KvCacheTransceiverV2) transceiver._send_sessions = sessions transceiver._send_reqs = reqs or {rid: _FakeRequest() for rid in sessions} + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + transceiver._wait_reqs = {} + transceiver._legacy_failed_sessions = set() + transceiver._context_cancelled_request_ids = [] + transceiver._async_ready_epoch = OrderedDict() + transceiver._async_ready_published = {} + transceiver._shutdown = False + transceiver._shutdown_complete = False + transceiver._shutdown_metadata_leases_complete = False + transceiver._shutdown_sessions_complete = False + transceiver._shutdown_consensus_complete = False + transceiver._shutdown_worker_complete = False + transceiver._shutdown_worker_event = None + transceiver._shutdown_deferred_errors = [] + transceiver._async_ready_metadata_leases = set() + transceiver._async_ready_idle_wakeup = threading.Event() transceiver._sender_future_timeout_ms = 123 + transceiver.kv_transfer_poll_interval_ms = 10 # Attributes read by check_context_transfer_status before it processes sessions. transceiver._ever_had_send_session = True transceiver._ctx_need_tp_sync = False @@ -107,9 +234,18 @@ def _make_transceiver( transceiver._transfer_worker = _FakeTransferWorker() transceiver._ctx_consensus = lambda local_ids: list(local_ids) transceiver._ctx_consensus_outcome = ( - lambda _to_process, cancelled, failed, completed, timed_out: ( + lambda _to_process, + _known_ids, + cancelled, + cancel_quiescent, + failed, + failed_quiescent, + completed, + timed_out: ( cancelled, + cancel_quiescent, failed, + failed_quiescent, completed, timed_out, ) @@ -117,6 +253,436 @@ def _make_transceiver( return transceiver +def _enable_fake_async_consensus( + transceiver: KvCacheTransceiverV2, + *, + terminal: bool = False, + peer_ready: bool = False, +) -> _FakeAsyncCoordinator: + coordinator = _FakeAsyncCoordinator() + transceiver._async_terminal_consensus_enabled = terminal + transceiver._async_peer_ready_consensus_enabled = peer_ready + transceiver._async_consensus = coordinator + transceiver._async_terminal_epoch = OrderedDict() + transceiver._async_terminal_published = {} + transceiver._async_terminal_commits = {} + transceiver._async_terminal_cancelled = {} + transceiver._async_ready_epoch = OrderedDict() + transceiver._async_ready_published = {} + transceiver._async_ready_prepared = {} + transceiver._async_ready_released = set() + transceiver._async_ready_activated = {} + transceiver._async_ready_acknowledged = set() + transceiver._async_ready_withdrawn = set() + transceiver._async_ready_aborted = {} + transceiver._async_ready_finalized_without_request = OrderedDict() + transceiver._async_ready_metadata_leases = set() + transceiver._async_consensus_counters = defaultdict(int) + transceiver._wait_reqs = {} + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + transceiver._dist = Mock(rank=0) + return coordinator + + +class _FakeStartupMpiDist: + def __init__(self, gather_result=None) -> None: + self.rank = 0 + self.gather_result = gather_result + self.descriptors: list = [] + self.pp_values: list = [] + + def allgather(self, descriptor): + self.descriptors.append(descriptor) + if self.gather_result is None: + return [descriptor, descriptor] + return self.gather_result(descriptor) + + def pp_allgather(self, value): + self.pp_values.append(value) + return [value, value] + + +def _make_startup_transceiver(dist: _FakeStartupMpiDist) -> KvCacheTransceiverV2: + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._dist = dist + transceiver._mapping = SimpleNamespace( + world_size=2, + tp_size=1, + pp_size=2, + cp_size=1, + enable_attention_dp=False, + pp_group=(0, 1), + ) + transceiver._transfer_worker = SimpleNamespace( + sender_endpoint="local-endpoint", + populate_instance_and_rank_info=Mock(), + ) + transceiver._kv_cache_manager = SimpleNamespace(pp_layers=[]) + transceiver._context_info_endpoint = "context-endpoint" + return transceiver + + +def test_async_startup_flag_off_preserves_legacy_endpoint_exchange(monkeypatch) -> None: + monkeypatch.setattr(transceiver_module, "MPIDist", _FakeStartupMpiDist) + dist = _FakeStartupMpiDist() + transceiver = _make_startup_transceiver(dist) + config = SimpleNamespace(backend="NIXL", transceiver_runtime="PYTHON") + + transceiver._init_async_consensus(config) + assert dist.descriptors == [] + + transceiver._exchange_rank_info() + + assert dist.descriptors == ["local-endpoint"] + transceiver._transfer_worker.populate_instance_and_rank_info.assert_called_once_with( + endpoints=["local-endpoint", "local-endpoint"], layer_num_per_pp=[0, 0] + ) + assert transceiver._async_consensus is None + + +def test_async_startup_rejects_same_version_flag_mismatch(monkeypatch) -> None: + def mismatch_flag(contribution): + tag, endpoint, descriptor = contribution + peer_descriptor = list(descriptor) + peer_descriptor[-2] = "0" + return [contribution, (tag, endpoint, tuple(peer_descriptor))] + + monkeypatch.setattr(transceiver_module, "MPIDist", _FakeStartupMpiDist) + monkeypatch.setenv(transceiver_module._ASYNC_TERMINAL_ENV, "1") + dist = _FakeStartupMpiDist(mismatch_flag) + transceiver = _make_startup_transceiver(dist) + config = SimpleNamespace(backend="NIXL", transceiver_runtime="PYTHON") + transceiver._init_async_consensus(config) + + with pytest.raises(RuntimeError, match="startup descriptor mismatch"): + transceiver._exchange_rank_info() + + assert len(dist.descriptors) == 1 + assert dist.pp_values == [0] + + +def test_async_startup_rejects_uniform_malformed_flag_after_exchange(monkeypatch) -> None: + monkeypatch.setattr(transceiver_module, "MPIDist", _FakeStartupMpiDist) + monkeypatch.setenv(transceiver_module._ASYNC_TERMINAL_ENV, "malformed") + dist = _FakeStartupMpiDist() + transceiver = _make_startup_transceiver(dist) + config = SimpleNamespace(backend="NIXL", transceiver_runtime="PYTHON") + transceiver._init_async_consensus(config) + + with pytest.raises(ValueError, match="must be 0 or 1"): + transceiver._exchange_rank_info() + + assert len(dist.descriptors) == 1 + assert dist.pp_values == [0] + + +def test_async_startup_rejects_mixed_legacy_worker_group(monkeypatch) -> None: + monkeypatch.setattr(transceiver_module, "MPIDist", _FakeStartupMpiDist) + monkeypatch.setenv(transceiver_module._ASYNC_TERMINAL_ENV, "1") + dist = _FakeStartupMpiDist(lambda contribution: [contribution, "legacy-peer-endpoint"]) + transceiver = _make_startup_transceiver(dist) + config = SimpleNamespace(backend="NIXL", transceiver_runtime="PYTHON") + transceiver._init_async_consensus(config) + + with pytest.raises(RuntimeError, match="same-version worker group"): + transceiver._exchange_rank_info() + + assert len(dist.descriptors) == 1 + assert dist.pp_values == [0] + + +def test_async_startup_rejects_unsupported_opt_in_after_exchange(monkeypatch) -> None: + monkeypatch.setattr(transceiver_module, "MPIDist", _FakeStartupMpiDist) + monkeypatch.setenv(transceiver_module._ASYNC_TERMINAL_ENV, "1") + dist = _FakeStartupMpiDist() + transceiver = _make_startup_transceiver(dist) + transceiver._mapping.tp_size = 2 + transceiver._mapping.pp_size = 1 + config = SimpleNamespace(backend="NIXL", transceiver_runtime="PYTHON") + transceiver._init_async_consensus(config) + + with pytest.raises(RuntimeError, match="currently requires"): + transceiver._exchange_rank_info() + + assert len(dist.descriptors) == 1 + assert dist.pp_values == [0] + + +@pytest.mark.parametrize( + "terminal_value,ready_value", + [("0", "0"), ("1", "0"), ("0", "1"), ("1", "1")], +) +def test_async_startup_supported_mode_matrix( + monkeypatch, terminal_value: str, ready_value: str +) -> None: + monkeypatch.setattr(transceiver_module, "MPIDist", _FakeStartupMpiDist) + monkeypatch.setenv(transceiver_module._ASYNC_TERMINAL_ENV, terminal_value) + monkeypatch.setenv(transceiver_module._ASYNC_PEER_READY_ENV, ready_value) + transport = object() + transport_constructor = Mock(return_value=transport) + coordinator = object() + coordinator_constructor = Mock(return_value=coordinator) + monkeypatch.setattr(transceiver_module, "MpiConsensusTransport", transport_constructor) + monkeypatch.setattr(transceiver_module, "AsyncConsensusCoordinator", coordinator_constructor) + dist = _FakeStartupMpiDist() + transceiver = _make_startup_transceiver(dist) + config = SimpleNamespace(backend="NIXL", transceiver_runtime="PYTHON") + + transceiver._init_async_consensus(config) + assert dist.descriptors == [] + transceiver._exchange_rank_info() + + assert transceiver._async_terminal_consensus_enabled is (terminal_value == "1") + assert transceiver._async_peer_ready_consensus_enabled is (ready_value == "1") + if terminal_value == "1" or ready_value == "1": + transport_constructor.assert_called_once_with((0, 1)) + coordinator_constructor.assert_called_once_with(transport, scheduling_rank=0) + assert transceiver._async_consensus is coordinator + else: + transport_constructor.assert_not_called() + coordinator_constructor.assert_not_called() + assert transceiver._async_consensus is None + + +def test_async_startup_closes_transport_when_coordinator_construction_fails( + monkeypatch, +) -> None: + monkeypatch.setattr(transceiver_module, "MPIDist", _FakeStartupMpiDist) + monkeypatch.setenv(transceiver_module._ASYNC_TERMINAL_ENV, "1") + transport = Mock() + monkeypatch.setattr( + transceiver_module, + "MpiConsensusTransport", + Mock(return_value=transport), + ) + monkeypatch.setattr( + transceiver_module, + "AsyncConsensusCoordinator", + Mock(side_effect=RuntimeError("coordinator construction failed")), + ) + dist = _FakeStartupMpiDist() + transceiver = _make_startup_transceiver(dist) + config = SimpleNamespace(backend="NIXL", transceiver_runtime="PYTHON") + + transceiver._init_async_consensus(config) + with pytest.raises(RuntimeError, match="coordinator construction failed"): + transceiver._exchange_rank_info() + + transport.close.assert_called_once_with(transceiver_module._CONSENSUS_STARTUP_CLOSE_TIMEOUT_S) + assert transceiver._async_consensus is None + + +def test_constructor_rolls_back_transfer_worker_when_consensus_startup_fails( + monkeypatch, +) -> None: + monkeypatch.setattr(transceiver_module, "MPIDist", _FakeStartupMpiDist) + monkeypatch.setenv(transceiver_module._ASYNC_TERMINAL_ENV, "1") + monkeypatch.setattr( + transceiver_module.torch.cuda, + "current_device", + Mock(return_value=0), + ) + monkeypatch.setattr( + KvCacheTransceiverV2, + "_check_compatible", + Mock(), + ) + monkeypatch.setattr( + KvCacheTransceiverV2, + "_init_sync_policy", + Mock(), + ) + monkeypatch.setattr( + KvCacheTransceiverV2, + "_broadcast_instance_name", + Mock(return_value="instance"), + ) + monkeypatch.setattr( + KvCacheTransceiverV2, + "_broadcast_context_endpoint", + Mock(return_value="context-endpoint"), + ) + monkeypatch.setattr( + transceiver_module, + "create_cache_reuse_adapter", + Mock(return_value=object()), + ) + worker = SimpleNamespace( + sender_endpoint="local-endpoint", + shutdown=Mock(return_value=None), + ) + monkeypatch.setattr(transceiver_module, "TransferWorker", Mock(return_value=worker)) + monkeypatch.setattr( + transceiver_module, + "MpiConsensusTransport", + Mock(side_effect=RuntimeError("transport construction failed")), + ) + mapping = SimpleNamespace( + world_size=2, + tp_size=1, + pp_size=2, + cp_size=1, + tp_rank=0, + enable_attention_dp=False, + pp_group=(0, 1), + ) + dist = _FakeStartupMpiDist() + kv_cache_manager = SimpleNamespace(max_batch_size=1, pp_layers=[]) + config = SimpleNamespace( + backend="NIXL", + transceiver_runtime="PYTHON", + kv_transfer_timeout_ms=1000, + kv_transfer_poll_interval_ms=10, + kv_transfer_sender_future_timeout_ms=1000, + kv_cache_bounce_size_mb=0, + ) + + with pytest.raises(RuntimeError, match="transport construction failed"): + KvCacheTransceiverV2(mapping, dist, kv_cache_manager, config) + + worker.shutdown.assert_called_once_with() + + +def test_startup_rollback_bounds_deferred_worker_wait(monkeypatch) -> None: + transceiver = _make_transceiver({}) + worker_event = threading.Event() + worker_event.wait = Mock(return_value=False) + transceiver.shutdown = Mock(return_value=worker_event) + + transceiver._rollback_failed_startup() + + worker_event.wait.assert_called_once_with(transceiver_module._STARTUP_ROLLBACK_TIMEOUT_S) + transceiver.shutdown.assert_called_once_with() + + +def test_shutdown_continues_after_session_close_failure_and_is_idempotent() -> None: + failed_session = _FakeSession(rid=41, wait_result=None) + failed_session.close = Mock(side_effect=RuntimeError("close failed")) + healthy_session = _FakeSession(rid=42, wait_result=None) + healthy_session.close = Mock() + transceiver = _make_transceiver( + {41: failed_session, 42: healthy_session}, + {41: _FakeRequest(request_id=41), 42: _FakeRequest(request_id=42)}, + ) + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + transceiver._async_consensus = Mock() + transceiver._async_consensus_counters = defaultdict(int) + transceiver._dist = SimpleNamespace(rank=0) + transceiver._transfer_worker.shutdown = Mock(return_value=None) + + with pytest.raises(RuntimeError, match="close failed"): + transceiver.shutdown() + + failed_session.close.assert_called_once_with() + healthy_session.close.assert_called_once_with() + transceiver._async_consensus.shutdown.assert_called_once_with() + transceiver._transfer_worker.shutdown.assert_called_once_with() + assert not transceiver._send_sessions + assert not transceiver._send_reqs + + transceiver.shutdown() + failed_session.close.assert_called_once_with() + transceiver._async_consensus.shutdown.assert_called_once_with() + transceiver._transfer_worker.shutdown.assert_called_once_with() + + +def test_shutdown_retries_only_incomplete_consensus_teardown() -> None: + transceiver = _make_transceiver({}) + transceiver._recv_sessions = {} + transceiver._recv_reqs = {} + transceiver._async_consensus = Mock() + transceiver._async_consensus.shutdown = Mock( + side_effect=[RuntimeError("peer shutdown timeout"), None] + ) + transceiver._async_consensus_counters = defaultdict(int) + transceiver._dist = SimpleNamespace(rank=0) + transceiver._transfer_worker.shutdown = Mock(return_value=None) + + with pytest.raises(RuntimeError, match="peer shutdown timeout"): + transceiver.shutdown() + + assert transceiver._shutdown + assert not transceiver._shutdown_complete + transceiver._async_consensus.shutdown.assert_called_once_with() + transceiver._transfer_worker.shutdown.assert_called_once_with() + + transceiver.shutdown() + + assert transceiver._shutdown_complete + assert transceiver._async_consensus.shutdown.call_count == 2 + # The worker completed on the first attempt and must not be torn down a + # second time while the delayed peer's consensus close is retried. + transceiver._transfer_worker.shutdown.assert_called_once_with() + + transceiver.shutdown() + assert transceiver._async_consensus.shutdown.call_count == 2 + + +def test_shutdown_retries_metadata_lease_before_stopping_worker() -> None: + transceiver = _make_transceiver({}) + transceiver._async_consensus = None + transceiver._async_ready_metadata_leases.add(46) + transceiver._transfer_worker.cancel_peer_result = False + transceiver._transfer_worker.shutdown = Mock(return_value=None) + + with pytest.raises(RuntimeError, match="metadata-only request teardown"): + transceiver.shutdown() + assert not transceiver._shutdown_metadata_leases_complete + assert not transceiver._shutdown_complete + transceiver._transfer_worker.shutdown.assert_not_called() + + transceiver._transfer_worker.cancel_peer_result = True + assert transceiver.shutdown() is None + assert transceiver._shutdown_metadata_leases_complete + assert transceiver._shutdown_complete + transceiver._transfer_worker.shutdown.assert_called_once_with() + + +def test_shutdown_propagates_deferred_worker_completion_before_marking_complete() -> None: + transceiver = _make_transceiver({}) + transceiver._async_consensus = None + completion = threading.Event() + transceiver._transfer_worker.shutdown = Mock(side_effect=[completion, None]) + + assert transceiver.shutdown() is completion + assert not transceiver._shutdown_worker_complete + assert not transceiver._shutdown_complete + + # Repeated calls while native teardown is live must propagate the same + # ownership barrier without invoking shutdown again. + assert transceiver.shutdown() is completion + transceiver._transfer_worker.shutdown.assert_called_once_with() + + completion.set() + assert transceiver.shutdown() is None + assert transceiver._shutdown_worker_complete + assert transceiver._shutdown_complete + assert transceiver._transfer_worker.shutdown.call_count == 2 + + +def test_shutdown_defers_primary_error_until_worker_completion() -> None: + session = _FakeSession(rid=43, wait_result=None) + session.close = Mock(side_effect=RuntimeError("session close failed")) + transceiver = _make_transceiver({43: session}) + transceiver._async_consensus = None + completion = threading.Event() + transceiver._transfer_worker.shutdown = Mock(side_effect=[completion, None]) + + # Native ownership is still live, so propagate the barrier instead of an + # error that could let the caller release registered memory too early. + assert transceiver.shutdown() is completion + + completion.set() + with pytest.raises(RuntimeError, match="session close failed"): + transceiver.shutdown() + + assert transceiver._shutdown_complete + session.close.assert_called_once_with() + assert transceiver._transfer_worker.shutdown.call_count == 2 + + def _make_tx_session( kv_tasks: list[_FakeTask], *, @@ -127,6 +693,11 @@ def _make_tx_session( session._timeout_s = 0.25 session._need_aux = need_aux session._terminal_status = None + session._terminal_snapshot = None + session._exception = None + session._outstanding_operations = 0 + session._sealed = False + session.lock = threading.Lock() session.receiver_ready = True session.kv_tasks = kv_tasks session.aux_task = aux_task @@ -152,6 +723,75 @@ def test_context_transfer_status_bounded_poll_keeps_not_ready_session_queued() - assert transceiver._transfer_worker.sweep_count == 1 +def test_legacy_context_cancel_retains_request_until_native_transfer_is_quiescent() -> None: + session = _FakeSession( + rid=14, + wait_result=None, + status=SessionStatus.CANCELLED, + has_transferring_tasks=True, + ) + req = _FakeRequest(request_id=14) + transceiver = _make_transceiver({14: session}, {14: req}) + + assert transceiver.check_context_transfer_status(0) == ([], []) + + assert transceiver._send_sessions[14] is session + assert transceiver._send_reqs[14] is req + assert not session.closed + + session._has_transferring_tasks = False + assert transceiver.check_context_transfer_status(0) == ([], []) + + assert 14 not in transceiver._send_sessions + assert 14 not in transceiver._send_reqs + assert session.closed + + +def test_legacy_peer_context_cancel_surfaces_quiescent_request_id() -> None: + session = _FakeSession(rid=16, wait_result=None) + req = _FakeRequest(request_id=16) + transceiver = _make_transceiver({16: session}, {16: req}) + transceiver._ctx_consensus_outcome = Mock(return_value=([16], [16], [], [], [], [])) + + # No local cancel_request() call: cancellation is learned from a peer's + # packed legacy outcome vote. + assert transceiver.check_context_transfer_status(0) == ([], []) + + assert session.cancel_calls == 1 + assert session.closed + assert 16 not in transceiver._send_sessions + assert 16 not in transceiver._send_reqs + assert transceiver.take_context_cancelled_request_ids() == [16] + assert transceiver.take_context_cancelled_request_ids() == [] + + +def test_legacy_context_failure_retains_request_until_sibling_operations_are_quiescent() -> None: + session = _FakeSession( + rid=15, + wait_result=WaitResult.FAILED, + status=SessionStatus.ERROR, + has_failed=True, + has_transferring_tasks=True, + ) + req = _FakeRequest(request_id=15) + transceiver = _make_transceiver({15: session}, {15: req}) + + assert transceiver.check_context_transfer_status(0) == ([], []) + + assert transceiver._send_sessions[15] is session + assert transceiver._send_reqs[15] is req + assert session.sealed + assert not session.closed + + session._has_transferring_tasks = False + assert transceiver.check_context_transfer_status(0) == ([], [15]) + + assert req.state == LlmRequestState.DISAGG_TRANS_ERROR + assert 15 not in transceiver._send_sessions + assert 15 not in transceiver._send_reqs + assert session.closed + + def test_context_transfer_status_block_all_uses_blocking_wait() -> None: session = _FakeSession(rid=12, wait_result=WaitResult.COMPLETED) req = _FakeRequest() @@ -191,6 +831,384 @@ def test_context_transfer_status_zero_budget_processes_task_level_failure() -> N assert 13 not in transceiver._send_reqs +def test_async_context_zero_budget_publishes_all_quiescent_terminal_sessions() -> None: + completed_session = _FakeSession( + rid=21, + wait_result=WaitResult.COMPLETED, + is_completed=True, + ) + transferring_session = _FakeSession( + rid=22, + wait_result=None, + has_transferring_tasks=True, + ) + transceiver = _make_transceiver({21: completed_session, 22: transferring_session}) + coordinator = _enable_fake_async_consensus(transceiver, terminal=True) + transceiver._ctx_consensus = Mock( + side_effect=AssertionError("legacy readiness collective must not run") + ) + transceiver._ctx_consensus_outcome = Mock( + side_effect=AssertionError("legacy outcome collective must not run") + ) + + completed, failed = transceiver.check_context_transfer_status(0) + + assert completed == [] + assert failed == [] + assert coordinator.terminal_votes == [(21, ConsensusOutcome.COMPLETED, 0)] + assert completed_session.sealed + assert completed_session.blocking_calls == [False] + assert not completed_session.closed + assert not transferring_session.blocking_calls + assert coordinator.poll_count == 2 + transceiver._ctx_consensus.assert_not_called() + transceiver._ctx_consensus_outcome.assert_not_called() + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.TERMINAL_COMMIT, + 21, + 0, + ConsensusOutcome.COMPLETED, + ) + ) + completed, failed = transceiver.check_context_transfer_status(0) + + assert completed == [21] + assert failed == [] + assert completed_session.closed + assert 21 not in transceiver._send_sessions + assert 22 in transceiver._send_sessions + + +def test_async_context_does_not_publish_if_work_starts_at_seal_boundary() -> None: + session = _FakeSession( + rid=24, + wait_result=WaitResult.FAILED, + has_failed=True, + seal_quiescent=False, + ) + transceiver = _make_transceiver({24: session}) + coordinator = _enable_fake_async_consensus(transceiver, terminal=True) + + assert transceiver.check_context_transfer_status(0) == ([], []) + assert session.sealed + assert coordinator.terminal_votes == [] + assert 24 in transceiver._send_sessions + + +def test_async_context_cancel_retains_session_until_authoritative_commit() -> None: + session = _FakeSession(rid=23, wait_result=WaitResult.FAILED) + req = _FakeRequest(request_id=23) + transceiver = _make_transceiver({23: session}, {23: req}) + coordinator = _enable_fake_async_consensus(transceiver, terminal=True) + + assert not transceiver.cancel_request(req) + assert coordinator.terminal_votes == [(23, ConsensusOutcome.CANCELLED, 0)] + assert 23 in transceiver._send_sessions + assert not session.closed + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.TERMINAL_COMMIT, + 23, + 0, + ConsensusOutcome.CANCELLED, + ) + ) + assert transceiver.cancel_request(req) + assert session.closed + assert 23 not in transceiver._send_sessions + # The same local retry consumed the tombstone, so a later status poll must + # not misclassify this already-terminated request as a peer cancellation. + assert transceiver.take_context_cancelled_request_ids() == [] + + +def test_async_cancelled_commit_status_poll_retains_owner_until_cancel_retry() -> None: + session = _FakeSession(rid=35, wait_result=WaitResult.FAILED) + req = _FakeRequest(request_id=35) + transceiver = _make_transceiver({35: session}, {35: req}) + coordinator = _enable_fake_async_consensus(transceiver, terminal=True) + + assert not transceiver.cancel_request(req) + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.TERMINAL_COMMIT, + 35, + 0, + ConsensusOutcome.CANCELLED, + ) + ) + + # The ordinary status poll wins the race with the executor's cancellation + # retry. Native resources close, but request-scoped ownership remains. + assert transceiver.check_context_transfer_status(0) == ([], []) + assert session.closed + assert 35 not in transceiver._send_sessions + assert transceiver.owns_request(req) + with pytest.raises(RuntimeError, match="before its asynchronous terminal cancellation"): + transceiver._get_or_create_send_session(_FakeRequest(request_id=35)) + + # The retry acknowledges exactly the cancelled transceiver leg and makes + # the next request-ID epoch reusable. + assert transceiver.cancel_request(req) + assert not transceiver.owns_request(req) + assert transceiver._async_terminal_epoch[35] == 1 + + +def test_async_peer_context_cancel_surfaces_id_without_local_retry() -> None: + session = _FakeSession(rid=38, wait_result=None) + req = _FakeRequest(request_id=38) + transceiver = _make_transceiver({38: session}, {38: req}) + coordinator = _enable_fake_async_consensus(transceiver, terminal=True) + transceiver._async_terminal_published[38] = 0 + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.TERMINAL_COMMIT, + 38, + 0, + ConsensusOutcome.CANCELLED, + ) + ) + + # No local cancel_request() call: the status poll must hand the peer's + # authoritative cancellation to PyExecutor while retaining its tombstone. + assert transceiver.check_context_transfer_status(0) == ([], []) + + assert session.closed + assert transceiver.owns_request(req) + assert transceiver.take_context_cancelled_request_ids() == [38] + assert transceiver.take_context_cancelled_request_ids() == [] + assert transceiver.cancel_request(req) + assert not transceiver.owns_request(req) + + +def test_async_context_cancel_does_not_replace_published_terminal_vote() -> None: + session = _FakeSession( + rid=25, + wait_result=WaitResult.COMPLETED, + is_completed=True, + ) + req = _FakeRequest(request_id=25) + transceiver = _make_transceiver({25: session}, {25: req}) + coordinator = _enable_fake_async_consensus(transceiver, terminal=True) + + assert transceiver.check_context_transfer_status(0) == ([], []) + assert coordinator.terminal_votes == [(25, ConsensusOutcome.COMPLETED, 0)] + + assert not transceiver.cancel_request(req) + assert session.cancel_calls == 0 + assert coordinator.terminal_votes == [(25, ConsensusOutcome.COMPLETED, 0)] + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.TERMINAL_COMMIT, + 25, + 0, + ConsensusOutcome.COMPLETED, + ) + ) + assert transceiver.cancel_request(req) + assert session.cancel_calls == 0 + assert session.closed + + +def test_async_context_positive_budget_waits_for_authoritative_commit() -> None: + session = _FakeSession( + rid=26, + wait_result=WaitResult.COMPLETED, + is_completed=True, + ) + transceiver = _make_transceiver({26: session}) + coordinator = _enable_fake_async_consensus(transceiver, terminal=True) + + def publish_commit_after_first_sweep(fake: _FakeAsyncCoordinator) -> None: + if fake.poll_count == 3 and fake.terminal_votes: + fake.events.append( + ConsensusEvent( + ConsensusEventKind.TERMINAL_COMMIT, + 26, + 0, + ConsensusOutcome.COMPLETED, + ) + ) + + coordinator.poll_hook = publish_commit_after_first_sweep + + completed, failed = transceiver.check_context_transfer_status(1) + + assert completed == [26] + assert failed == [] + assert coordinator.poll_count >= 3 + assert session.blocking_calls == [False] + + +def test_async_context_block_all_drains_entry_snapshot_through_commits() -> None: + sessions = { + rid: _FakeSession( + rid=rid, + wait_result=WaitResult.COMPLETED, + is_completed=True, + ) + for rid in (27, 28) + } + session_objects = list(sessions.values()) + transceiver = _make_transceiver(sessions) + coordinator = _enable_fake_async_consensus(transceiver, terminal=True) + + def publish_commits_after_first_sweep(fake: _FakeAsyncCoordinator) -> None: + if fake.poll_count != 3: + return + for rid, _outcome, epoch in fake.terminal_votes: + fake.events.append( + ConsensusEvent( + ConsensusEventKind.TERMINAL_COMMIT, + rid, + epoch, + ConsensusOutcome.COMPLETED, + ) + ) + + coordinator.poll_hook = publish_commits_after_first_sweep + + completed, failed = transceiver.check_context_transfer_status(None) + + assert completed == [27, 28] + assert failed == [] + assert not transceiver._send_sessions + assert all(session.blocking_calls == [False] for session in session_objects) + + +def test_async_context_block_all_stops_at_existing_sender_timeout() -> None: + session = _FakeSession(rid=30, wait_result=None) + transceiver = _make_transceiver({30: session}) + _enable_fake_async_consensus(transceiver, terminal=True) + transceiver._sender_future_timeout_ms = 1 + + assert transceiver.check_context_transfer_status(None) == ([], []) + + assert 30 in transceiver._send_sessions + assert session.blocking_calls + assert all(not blocking for blocking in session.blocking_calls) + + +def test_async_context_uses_native_atomic_terminal_snapshot_when_available() -> None: + session = _FakeSession(rid=29, wait_result=None) + session.seal_and_snapshot_terminal = Mock( # type: ignore[attr-defined] + return_value=SessionStatus.KV_TRANSFERRED + ) + session.seal_and_check_quiescent = Mock( # type: ignore[method-assign] + side_effect=AssertionError("compatibility seal must not run") + ) + transceiver = _make_transceiver({29: session}) + coordinator = _enable_fake_async_consensus(transceiver, terminal=True) + + assert transceiver.check_context_transfer_status(0) == ([], []) + + assert coordinator.terminal_votes == [(29, ConsensusOutcome.COMPLETED, 0)] + session.seal_and_snapshot_terminal.assert_called_once_with() # type: ignore[attr-defined] + + +def test_respond_pre_cancel_keeps_paired_ownership_without_dispatch() -> None: + transceiver = _make_transceiver({}) + session = _FakeSession( + rid=41, + wait_result=WaitResult.FAILED, + status=SessionStatus.CANCELLED, + ) + session.send = Mock() # type: ignore[attr-defined] + transceiver._transfer_worker.tx_session = session + transceiver._create_kv_slice = Mock(side_effect=AssertionError("must not build a slice")) + req = _FakeRequest(request_id=41) + + transceiver.respond_and_send_async(req) + + assert transceiver._send_sessions[41] is session + assert transceiver._send_reqs[41] is req + session.send.assert_not_called() # type: ignore[attr-defined] + + +def test_respond_cancel_between_kv_and_aux_keeps_paired_ownership() -> None: + transceiver = _make_transceiver({}) + session = _FakeSession(rid=42, wait_result=None) + session.send = Mock() # type: ignore[attr-defined] + session.pack_aux = Mock( # type: ignore[attr-defined] + side_effect=lambda _req: session.cancel() + ) + session.send_aux = Mock( # type: ignore[attr-defined] + side_effect=RuntimeError("session sealed by cancellation") + ) + transceiver._transfer_worker.tx_session = session + transceiver._create_kv_slice = Mock(return_value=Mock()) + transceiver._need_aux_transfer = lambda _req: True + transceiver._dp_rank = 0 + transceiver._context_info_endpoint = "ctx" + req = _FakeRequest(request_id=42) + + transceiver.respond_and_send_async(req) + + assert transceiver._send_sessions[42] is session + assert transceiver._send_reqs[42] is req + session.send.assert_called_once() # type: ignore[attr-defined] + session.send_aux.assert_called_once() # type: ignore[attr-defined] + + +def test_respond_dispatch_error_preserves_paired_ownership_before_reraising() -> None: + transceiver = _make_transceiver({}) + session = _FakeSession(rid=43, wait_result=None) + session.send = Mock(side_effect=RuntimeError("dispatch failed")) # type: ignore[attr-defined] + transceiver._transfer_worker.tx_session = session + transceiver._create_kv_slice = Mock(return_value=Mock()) + req = _FakeRequest(request_id=43) + + with pytest.raises(RuntimeError, match="dispatch failed"): + transceiver.respond_and_send_async(req) + + assert transceiver._send_sessions[43] is session + assert transceiver._send_reqs[43] is req + + +def test_receive_pre_cancel_keeps_paired_ownership_without_dispatch() -> None: + transceiver = _make_transceiver({}) + session = _FakeSession( + rid=44, + wait_result=WaitResult.FAILED, + status=SessionStatus.CANCELLED, + ) + session.receive = Mock() # type: ignore[attr-defined] + transceiver._transfer_worker.rx_session = session + transceiver._create_kv_slice = Mock(side_effect=AssertionError("must not build a slice")) + req = _FakeRequest(request_id=44) + + transceiver.request_and_receive_async(req) + + assert transceiver._recv_sessions[44] is session + assert transceiver._recv_reqs[44] is req + session.receive.assert_not_called() # type: ignore[attr-defined] + + +def test_receive_cancel_during_dispatch_keeps_paired_ownership() -> None: + transceiver = _make_transceiver({}) + session = _FakeSession(rid=45, wait_result=None) + + def cancel_and_raise(_slice) -> None: + session.cancel() + raise RuntimeError("session sealed by cancellation") + + session.receive = Mock(side_effect=cancel_and_raise) # type: ignore[attr-defined] + transceiver._transfer_worker.rx_session = session + transceiver._create_kv_slice = Mock(return_value=Mock()) + transceiver._slice_num_bytes = Mock(return_value=0) + transceiver._kv_size_rank_factor = 1 + req = _FakeRequest(request_id=45) + + transceiver.request_and_receive_async(req) + + assert transceiver._recv_sessions[45] is session + assert transceiver._recv_reqs[45] is req + session.receive.assert_called_once() # type: ignore[attr-defined] + + def test_context_transfer_status_skips_consensus_when_never_sent() -> None: # A worker that never sends skips the ctx consensus even when TP sync would need it, but still # sweeps so nothing leaks. @@ -236,7 +1254,7 @@ def test_gen_transfer_status_enters_consensus_when_sync_required() -> None: transceiver._recv_reqs = {} transceiver._gen_consensus = Mock(return_value=[]) transceiver._build_to_process = Mock(return_value=[]) - transceiver._gen_consensus_outcome = Mock(return_value=([], [], [])) + transceiver._gen_consensus_outcome = Mock(return_value=([], [], [], [], [])) transceiver._close_failed_sessions = Mock() completed, failed, cancelled = transceiver.check_gen_transfer_status(at_least_request_num=0) @@ -247,27 +1265,108 @@ def test_gen_transfer_status_enters_consensus_when_sync_required() -> None: transceiver._gen_consensus.assert_called_once_with([]) +def test_legacy_gen_cancel_retains_request_until_native_cancel_ack_is_quiescent() -> None: + session = _FakeSession( + rid=51, + wait_result=None, + status=SessionStatus.CANCELLED, + has_transferring_tasks=True, + ) + req = _FakeRequest(request_id=51) + transceiver = _make_transceiver({}) + transceiver._ever_had_recv_session = True + transceiver._gen_need_sync = False + transceiver._gen_allgather = Mock() + transceiver._gen_consensus = lambda local_ids: list(local_ids) + transceiver._recv_sessions = {51: session} + transceiver._recv_reqs = {51: req} + + assert transceiver.check_gen_transfer_status(0) == ([], [], []) + + assert transceiver._recv_sessions[51] is session + assert transceiver._recv_reqs[51] is req + assert not session.closed + + # RxSession.has_transferring_tasks() includes the sender cancellation ACK, + # so this transition models both the active write and its ACK draining. + session._has_transferring_tasks = False + completed, failed, cancelled = transceiver.check_gen_transfer_status(0) + + assert completed == [] + assert failed == [] + assert cancelled == [req] + assert 51 not in transceiver._recv_sessions + assert 51 not in transceiver._recv_reqs + assert session.closed + + +def test_legacy_gen_failure_retains_request_until_sibling_operations_are_quiescent() -> None: + session = _FakeSession( + rid=52, + wait_result=WaitResult.FAILED, + status=SessionStatus.ERROR, + has_failed=True, + has_transferring_tasks=True, + ) + req = _FakeRequest(request_id=52) + transceiver = _make_transceiver({}) + transceiver._ever_had_recv_session = True + transceiver._gen_need_sync = False + transceiver._gen_allgather = Mock() + transceiver._gen_consensus = lambda local_ids: list(local_ids) + transceiver._recv_sessions = {52: session} + transceiver._recv_reqs = {52: req} + transceiver._dist = SimpleNamespace(rank=0) + + assert transceiver.check_gen_transfer_status(0) == ([], [], []) + + assert transceiver._recv_sessions[52] is session + assert transceiver._recv_reqs[52] is req + assert session.sealed + assert not session.closed + + session._has_transferring_tasks = False + assert transceiver.check_gen_transfer_status(0) == ([], [52], []) + + assert req.state == LlmRequestState.DISAGG_TRANS_ERROR + assert 52 not in transceiver._recv_sessions + assert 52 not in transceiver._recv_reqs + assert session.closed + + def test_consensus_outcome_uses_single_batched_allgather() -> None: - # The cancelled/failed/completed id lists are exchanged with ONE allgather - # (packed as a list-of-lists) instead of three; verify a single call and that - # union (cancelled/failed) + intersection (completed) semantics are preserved. + # Decision, quiescence, failure, and completion are exchanged with ONE + # allgather. Verify union (cancel/fail) and intersection + # (cancel-quiescence/complete) semantics without an extra rendezvous. transceiver = object.__new__(KvCacheTransceiverV2) calls: list = [] def fake_allgather(payload): calls.append(payload) - # rank0 = this rank's [cancelled, failed, completed]; rank1 = a peer rank. - return [payload, [[], [99], [7, 8]]] + # rank0 = this rank's payload; rank1 = a peer rank. + return [payload, [[], [1], [99], [2, 99], [7, 8]]] to_process = [1, 2, 7, 8, 99] - new_cancelled, new_failed, new_completed = transceiver._consensus_outcome( - to_process, [1], [2], [7], fake_allgather, True + new_cancelled, reclaimable_cancelled, new_failed, reclaimable_failed, new_completed = ( + transceiver._consensus_outcome( + to_process, + to_process, + [1], + [1], + [2], + [2], + [7], + fake_allgather, + True, + ) ) assert len(calls) == 1 # batched: a single allgather, not three - assert calls[0] == [[1], [2], [7]] + assert calls[0] == [[1], [1], [2], [2], [7]] assert new_cancelled == [1] # union of cancelled across ranks + assert reclaimable_cancelled == [1] # quiescent only because both ranks ACKed assert new_failed == [2, 99] # union of failed across ranks + assert reclaimable_failed == [2] # 99 lacks this rank's quiescence ACK assert new_completed == [7] # intersection only (8 is completed on the peer only) @@ -313,6 +1412,33 @@ def test_ctx_consensus_fastpath_skips_when_idle(monkeypatch) -> None: transceiver._ctx_consensus.assert_called_once() +def test_consensus_outcome_defers_cancel_reclamation_until_every_rank_is_quiescent() -> None: + transceiver = object.__new__(KvCacheTransceiverV2) + + def fake_allgather(payload): + # The peer has observed the same cancellation decision but still has + # an active native transfer, so it omits the quiescence ACK. + return [payload, [[61], [], [], [], []]] + + cancelled, reclaimable, failed, reclaimable_failed, completed = transceiver._consensus_outcome( + [61], + [61], + [61], + [61], + [], + [], + [], + fake_allgather, + True, + ) + + assert cancelled == [61] + assert reclaimable == [] + assert failed == [] + assert reclaimable_failed == [] + assert completed == [] + + def test_tx_session_wait_complete_defaults_to_blocking() -> None: task = _FakeTask(TaskStatus.INIT, wait_result=False) session = _make_tx_session([task]) @@ -352,7 +1478,7 @@ def test_check_context_runs_consensus_after_a_send() -> None: transceiver._ever_had_send_session = True transceiver._ctx_need_tp_sync = True transceiver._ctx_consensus = Mock(return_value=[]) - transceiver._ctx_consensus_outcome = Mock(return_value=([], [], [], [])) + transceiver._ctx_consensus_outcome = Mock(return_value=([], [], [], [], [], [])) transceiver.check_context_transfer_status(0) transceiver._ctx_consensus.assert_called_once() @@ -367,3 +1493,616 @@ def test_prepare_context_requests_skips_consensus_when_nothing_waiting() -> None transceiver.prepare_context_requests([]) transceiver._ctx_consensus.assert_not_called() + + +def test_context_activation_digest_is_common_to_legacy_and_authoritative_paths( + monkeypatch, +) -> None: + monkeypatch.setenv(transceiver_module._CONTEXT_ACTIVATION_DIGEST_ENV, "1") + requests = [_FakeRequest(request_id=11), _FakeRequest(request_id=22)] + + legacy = _make_transceiver({}) + legacy._init_context_activation_digest() + legacy._wait_reqs = {request.request_id: request for request in requests} + legacy._transfer_worker.ready_request_ids.update(legacy._wait_reqs) + legacy.prepare_context_requests([]) + + authoritative = _make_transceiver({}) + _enable_fake_async_consensus(authoritative, peer_ready=True) + authoritative._init_context_activation_digest() + authoritative._async_ready_published = {request.request_id: 0 for request in requests} + authoritative._async_ready_prepared = {(request.request_id, 0): request for request in requests} + authoritative._async_ready_released = {(request.request_id, 0) for request in requests} + authoritative.activate_context_requests_for_schedule(requests) + + assert legacy._context_activation_count == 2 + assert authoritative._context_activation_count == 2 + assert ( + legacy._context_activation_digest.hexdigest() + == authoritative._context_activation_digest.hexdigest() + == "f5e6334494e9cf39f8dbc5e0b404ddc0d2d0491dc8d469705b228e7a8e1b4aa9" + ) + + reversed_order = _make_transceiver({}) + reversed_order._init_context_activation_digest() + reversed_order._record_context_activation_ids([22, 11]) + assert ( + reversed_order._context_activation_digest.hexdigest() + == "ae0d5811960bd6babae707df768bf641bb04a9eb2d628cc95f2194c39aef55cd" + ) + + +def test_context_activation_digest_is_default_off(monkeypatch) -> None: + transceiver = _make_transceiver({}) + transceiver._init_context_activation_digest() + transceiver._record_context_activation_ids([11, 22]) + info = Mock() + monkeypatch.setattr(transceiver_module.logger, "info", info) + + transceiver._log_context_activation_digest() + + assert transceiver._context_activation_digest is None + assert transceiver._context_activation_count == 0 + info.assert_not_called() + + +def test_context_activation_digest_logs_once_after_shutdown(monkeypatch) -> None: + monkeypatch.setenv(transceiver_module._CONTEXT_ACTIVATION_DIGEST_ENV, "1") + transceiver = _make_transceiver({}) + transceiver._init_context_activation_digest() + transceiver._record_context_activation_ids([11, 22]) + transceiver._async_consensus = None + transceiver._dist = SimpleNamespace(rank=3) + transceiver._transfer_worker.shutdown = Mock(return_value=None) + info = Mock() + monkeypatch.setattr(transceiver_module.logger, "info", info) + + transceiver.shutdown() + transceiver.shutdown() + + info.assert_called_once_with( + "PYTHON_CONTEXT_ACTIVATION_SEQUENCE " + "rank=3 count=2 " + "digest=f5e6334494e9cf39f8dbc5e0b404ddc0d2d0491dc8d469705b228e7a8e1b4aa9 " + "algorithm=sha256-length-prefixed-decimal-v1" + ) + + +def test_async_peer_ready_pins_metadata_while_request_is_pre_active() -> None: + transceiver = _make_transceiver({}) + _enable_fake_async_consensus(transceiver, peer_ready=True) + req = _FakeRequest(request_id=30) + + transceiver.prepare_context_requests([req]) + + assert transceiver.supports_pre_active_context_requests() + assert transceiver._transfer_worker.pinned_request_ids == {30} + assert transceiver._async_ready_metadata_leases == {30} + assert transceiver.owns_request(req) + + +def test_async_peer_ready_activates_only_from_authoritative_schedule() -> None: + transceiver = _make_transceiver({}) + coordinator = _enable_fake_async_consensus(transceiver, peer_ready=True) + req = _FakeRequest(request_id=31) + transceiver._wait_reqs[31] = req + transceiver._transfer_worker.ready_request_ids.add(31) + + transceiver._prepare_context_requests_async() + assert coordinator.ready_votes == [(31, 0)] + assert req.state is None + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_PREPARE, + 31, + 0, + ConsensusOutcome.READY, + ) + ) + transceiver._progress_async_consensus() + assert req.state is None + assert 31 not in transceiver._wait_reqs + assert coordinator.ready_acks == [(31, 0)] + assert not transceiver.is_context_request_ready_for_activation(req) + + # Omitting the request from rank zero's schedule leaves the lease hidden. + transceiver.activate_context_requests_for_schedule([]) + assert req.state is None + assert coordinator.ready_activation_acks == [] + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_RELEASE, + 31, + 0, + ConsensusOutcome.READY, + ) + ) + transceiver._progress_async_consensus() + assert req.state == LlmRequestState.CONTEXT_INIT + assert (31, 0) in transceiver._async_ready_prepared + assert transceiver.is_context_request_ready_for_activation(req) + + transceiver.activate_context_requests_for_schedule([req]) + assert coordinator.ready_activation_acks == [(31, 0)] + assert not transceiver._async_ready_prepared + assert (31, 0) in transceiver._async_ready_activated + assert 31 in transceiver._async_ready_published + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_COMPLETE, + 31, + 0, + ConsensusOutcome.READY, + ) + ) + transceiver._progress_async_consensus() + assert not transceiver._async_ready_activated + assert 31 not in transceiver._async_ready_published + + +def test_async_peer_ready_activation_validates_batch_atomically() -> None: + transceiver = _make_transceiver({}) + coordinator = _enable_fake_async_consensus(transceiver, peer_ready=True) + first = _FakeRequest(request_id=41) + second = _FakeRequest(request_id=42) + wrong_second = _FakeRequest(request_id=42) + transceiver._async_ready_published.update({41: 0, 42: 0}) + transceiver._async_ready_prepared.update( + { + (41, 0): first, + (42, 0): second, + } + ) + transceiver._async_ready_released.update({(41, 0), (42, 0)}) + + with pytest.raises(RuntimeError, match="does not match PREPARE"): + transceiver.activate_context_requests_for_schedule([first, wrong_second]) + + assert transceiver._async_ready_prepared == { + (41, 0): first, + (42, 0): second, + } + assert transceiver._async_ready_activated == {} + assert coordinator.ready_activation_acks == [] + + +def test_async_peer_ready_completion_retains_metadata_ownership_until_session() -> None: + transceiver = _make_transceiver({}) + coordinator = _enable_fake_async_consensus(transceiver, peer_ready=True) + req = _FakeRequest(request_id=43) + transceiver._transfer_worker.ready_request_ids.add(43) + transceiver.prepare_context_requests([req]) + coordinator.events.extend( + [ + ConsensusEvent(ConsensusEventKind.READY_PREPARE, 43, 0, ConsensusOutcome.READY), + ConsensusEvent(ConsensusEventKind.READY_RELEASE, 43, 0, ConsensusOutcome.READY), + ] + ) + transceiver._progress_async_consensus() + transceiver.activate_context_requests_for_schedule([req]) + coordinator.events.append( + ConsensusEvent(ConsensusEventKind.READY_COMPLETE, 43, 0, ConsensusOutcome.READY) + ) + transceiver._progress_async_consensus() + + assert transceiver.owns_request(req) + assert transceiver.cancel_request(req) + assert not transceiver.owns_request(req) + assert 43 in transceiver._transfer_worker.cancelled_peer_request_ids + + +def test_async_peer_ready_metadata_ownership_transfers_to_tx_session() -> None: + transceiver = _make_transceiver({}) + _enable_fake_async_consensus(transceiver, peer_ready=True) + req = _FakeRequest(request_id=44) + session = Mock() + transceiver._transfer_worker.tx_session = session + transceiver._async_ready_metadata_leases.add(44) + + assert transceiver._get_or_create_send_session(req) is session + assert 44 not in transceiver._async_ready_metadata_leases + assert transceiver._send_sessions[44] is session + + +def test_async_peer_ready_cancel_retains_ownership_if_native_handoff_won() -> None: + transceiver = _make_transceiver({}) + _enable_fake_async_consensus(transceiver, peer_ready=True) + req = _FakeRequest(request_id=45) + transceiver._async_ready_metadata_leases.add(45) + transceiver._transfer_worker.cancel_peer_result = False + + assert not transceiver.cancel_request(req) + assert transceiver.owns_request(req) + assert 45 in transceiver._async_ready_metadata_leases + + +def test_async_peer_ready_wait_uses_bounded_interruptible_backoff(monkeypatch) -> None: + transceiver = _make_transceiver({}) + _enable_fake_async_consensus(transceiver, peer_ready=True) + transceiver._wait_reqs[37] = _FakeRequest(request_id=37) + transceiver.kv_transfer_poll_interval_ms = 5000 + wait = Mock(return_value=False) + monkeypatch.setattr(transceiver._async_ready_idle_wakeup, "wait", wait) + + transceiver._prepare_context_requests_async() + + wait.assert_called_once_with(transceiver_module._ASYNC_READY_MAX_IDLE_SLEEP_S) + assert 37 not in transceiver._async_ready_published + + +def test_async_peer_ready_progress_does_not_backoff(monkeypatch) -> None: + transceiver = _make_transceiver({}) + coordinator = _enable_fake_async_consensus(transceiver, peer_ready=True) + req = _FakeRequest(request_id=38) + transceiver._wait_reqs[38] = req + transceiver._async_ready_published[38] = 0 + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_PREPARE, + 38, + 0, + ConsensusOutcome.READY, + ) + ) + wait = Mock(return_value=False) + monkeypatch.setattr(transceiver._async_ready_idle_wakeup, "wait", wait) + + transceiver._prepare_context_requests_async() + + wait.assert_not_called() + assert (38, 0) in transceiver._async_ready_prepared + + +def test_async_peer_ready_shutdown_skips_backoff(monkeypatch) -> None: + transceiver = _make_transceiver({}) + _enable_fake_async_consensus(transceiver, peer_ready=True) + transceiver._wait_reqs[39] = _FakeRequest(request_id=39) + transceiver._shutdown = True + wait = Mock(return_value=False) + monkeypatch.setattr(transceiver._async_ready_idle_wakeup, "wait", wait) + + transceiver._prepare_context_requests_async() + + wait.assert_not_called() + + +def test_estimation_transceiver_never_publishes_endpoint_metadata() -> None: + transceiver = _make_transceiver({}) + transceiver._publish_disaggregated_params = False + transceiver._mapping = SimpleNamespace(enable_attention_dp=False) + transceiver._dp_rank = 0 + transceiver._context_info_endpoint = "tcp://profiling:1234" + transceiver._instance_name = "profiling-generation" + transceiver._async_peer_ready_consensus_enabled = False + + assert transceiver.get_disaggregated_params() == {} + + transceiver._publish_disaggregated_params = True + assert transceiver.get_disaggregated_params() == { + "ctx_dp_rank": 0, + "ctx_info_endpoint": ["tcp://profiling:1234"], + } + + transceiver._async_peer_ready_consensus_enabled = True + assert transceiver.get_disaggregated_params() == { + "ctx_dp_rank": 0, + "ctx_info_endpoint": ["tcp://profiling:1234"], + "ctx_endpoint_generation": "profiling-generation", + } + + +def test_async_peer_ready_follower_stays_hidden_until_schedule_arrives() -> None: + transceiver = _make_transceiver({}) + coordinator = _enable_fake_async_consensus(transceiver, peer_ready=True) + transceiver._dist.rank = 1 + req = _FakeRequest(request_id=35) + transceiver._wait_reqs[35] = req + transceiver._async_ready_published[35] = 0 + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_PREPARE, + 35, + 0, + ConsensusOutcome.READY, + ) + ) + transceiver._progress_async_consensus() + + assert req.state is None + assert transceiver.is_context_request_ready_for_activation(req) + assert not transceiver.is_context_request_ready_for_activation(_FakeRequest(request_id=99)) + transceiver.activate_context_requests_for_schedule([]) + assert req.state is None + + transceiver.activate_context_requests_for_schedule([req]) + assert req.state == LlmRequestState.CONTEXT_INIT + assert coordinator.ready_activation_acks == [(35, 0)] + + +def test_async_peer_ready_cancel_withdraws_and_tombstones_request() -> None: + transceiver = _make_transceiver({}) + coordinator = _enable_fake_async_consensus(transceiver, peer_ready=True) + req = _FakeRequest(request_id=32) + transceiver._transfer_worker.ready_request_ids.add(32) + transceiver.prepare_context_requests([req]) + + assert not transceiver.cancel_request(req) + assert coordinator.ready_withdrawals == [(32, 0)] + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_ABORT, + 32, + 0, + ConsensusOutcome.WITHDRAWN, + ) + ) + transceiver._progress_async_consensus() + assert coordinator.ready_abort_acks == [(32, 0)] + assert not transceiver.cancel_request(req) + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_ABORT_FINALIZE, + 32, + 0, + ConsensusOutcome.WITHDRAWN, + ) + ) + transceiver._progress_async_consensus() + assert transceiver.cancel_request(req) + assert 32 in transceiver._transfer_worker.cancelled_peer_request_ids + assert 32 not in transceiver._async_ready_metadata_leases + + transceiver.prepare_context_requests([req]) + assert coordinator.ready_votes == [(32, 0)] + assert 32 not in transceiver._wait_reqs + + +def test_async_peer_ready_cancel_before_local_vote_waits_for_abort_finalize() -> None: + transceiver = _make_transceiver({}) + coordinator = _enable_fake_async_consensus(transceiver, peer_ready=True) + req = _FakeRequest(request_id=36) + transceiver._wait_reqs[36] = req + + # Local peer metadata is not ready, but another rank may already have + # voted. Join the current epoch with a withdrawal instead of freeing early. + assert not transceiver.cancel_request(req) + assert coordinator.ready_votes == [] + assert coordinator.ready_withdrawals == [(36, 0)] + assert transceiver.owns_request(req) + assert 36 in transceiver._wait_reqs + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_ABORT, + 36, + 0, + ConsensusOutcome.WITHDRAWN, + ) + ) + transceiver._progress_async_consensus() + assert coordinator.ready_abort_acks == [(36, 0)] + assert not transceiver.cancel_request(req) + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_ABORT_FINALIZE, + 36, + 0, + ConsensusOutcome.WITHDRAWN, + ) + ) + transceiver._progress_async_consensus() + assert transceiver.cancel_request(req) + assert not transceiver.owns_request(req) + assert transceiver._async_ready_epoch[36] == 1 + + +def test_async_peer_ready_abort_before_local_request_binds_without_fail_stop() -> None: + transceiver = _make_transceiver({}) + coordinator = _enable_fake_async_consensus(transceiver, peer_ready=True) + req = _FakeRequest(request_id=40) + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_ABORT, + 40, + 0, + ConsensusOutcome.WITHDRAWN, + ) + ) + transceiver._progress_async_consensus() + + assert coordinator.ready_abort_acks == [(40, 0)] + assert transceiver._async_ready_aborted[(40, 0)] is None + + transceiver.prepare_context_requests([req]) + + assert transceiver._async_ready_aborted[(40, 0)] is req + assert req.state == LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER + assert req._trtllm_async_ready_cancelled_epoch == 0 + assert 40 not in transceiver._wait_reqs + assert coordinator.ready_votes == [] + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_ABORT_FINALIZE, + 40, + 0, + ConsensusOutcome.WITHDRAWN, + ) + ) + transceiver._progress_async_consensus() + assert transceiver.cancel_request(req) + + +def test_async_peer_ready_finalized_abort_excludes_late_request() -> None: + transceiver = _make_transceiver({}) + coordinator = _enable_fake_async_consensus(transceiver, peer_ready=True) + req = _FakeRequest(request_id=41) + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_ABORT, + 41, + 0, + ConsensusOutcome.WITHDRAWN, + ) + ) + transceiver._progress_async_consensus() + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_ABORT_FINALIZE, + 41, + 0, + ConsensusOutcome.WITHDRAWN, + ) + ) + transceiver._progress_async_consensus() + + assert transceiver._async_ready_finalized_without_request[41] == 0 + assert transceiver._async_ready_epoch[41] == 1 + + transceiver.prepare_context_requests([req]) + + assert 41 not in transceiver._async_ready_finalized_without_request + assert req.state == LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER + assert req._trtllm_async_ready_cancelled_epoch == 0 + assert 41 not in transceiver._wait_reqs + assert coordinator.ready_votes == [] + + +def test_known_cancel_withdraws_before_readiness_poll_or_vote() -> None: + transceiver = _make_transceiver({}) + coordinator = _enable_fake_async_consensus(transceiver, peer_ready=True) + req = _FakeRequest(request_id=37) + transceiver._wait_reqs[37] = req + transceiver._transfer_worker.ready_request_ids.add(37) + + transceiver.exclude_context_requests_from_readiness([req]) + + assert coordinator.poll_count == 0 + assert coordinator.ready_votes == [] + assert coordinator.ready_withdrawals == [(37, 0)] + assert req.state == LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER + assert transceiver.owns_request(req) + + # The immediately following readiness progression observes the + # cancellation marker and cannot publish a contradictory READY vote. + transceiver.prepare_context_requests([req]) + assert coordinator.ready_votes == [] + + +def test_cancel_before_readiness_tracks_and_retires_native_metadata() -> None: + transceiver = _make_transceiver({}) + coordinator = _enable_fake_async_consensus(transceiver, peer_ready=True) + req = _FakeRequest(request_id=42) + + transceiver.exclude_context_requests_from_readiness([req]) + + assert coordinator.ready_votes == [] + assert coordinator.ready_withdrawals == [] + assert 42 in transceiver._transfer_worker.pinned_request_ids + assert 42 in transceiver._async_ready_metadata_leases + assert transceiver.owns_request(req) + + assert transceiver.cancel_request(req) + assert 42 in transceiver._transfer_worker.cancelled_peer_request_ids + assert 42 not in transceiver._async_ready_metadata_leases + + +def test_default_off_cancel_removes_existing_legacy_readiness_waiter() -> None: + transceiver = _make_transceiver({}) + transceiver._async_peer_ready_consensus_enabled = False + req = _FakeRequest(request_id=39) + transceiver._wait_reqs[39] = req + transceiver._transfer_worker.ready_request_ids.add(39) + + transceiver.exclude_context_requests_from_readiness([req]) + + assert 39 not in transceiver._wait_reqs + assert req.state == LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER + transceiver.prepare_context_requests([]) + assert 39 not in transceiver._wait_reqs + assert req.state == LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER + + +def test_async_peer_ready_abort_does_not_requeue_prepared_request() -> None: + transceiver = _make_transceiver({}) + coordinator = _enable_fake_async_consensus(transceiver, peer_ready=True) + req = _FakeRequest(request_id=33) + transceiver._wait_reqs[33] = req + transceiver._async_ready_published[33] = 0 + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_PREPARE, + 33, + 0, + ConsensusOutcome.READY, + ) + ) + transceiver._progress_async_consensus() + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_ABORT, + 33, + 0, + ConsensusOutcome.WITHDRAWN, + ) + ) + transceiver._progress_async_consensus() + + assert req.state == LlmRequestState.DISAGG_CONTEXT_WAIT_SCHEDULER + assert 33 not in transceiver._wait_reqs + assert not transceiver._async_ready_prepared + assert req._trtllm_async_ready_cancelled_epoch == 0 + assert coordinator.ready_abort_acks == [(33, 0)] + + +def test_async_peer_ready_cancel_after_ack_waits_for_completion() -> None: + transceiver = _make_transceiver({}) + coordinator = _enable_fake_async_consensus(transceiver, peer_ready=True) + req = _FakeRequest(request_id=34) + transceiver._wait_reqs[34] = req + transceiver._async_ready_published[34] = 0 + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_PREPARE, + 34, + 0, + ConsensusOutcome.READY, + ) + ) + transceiver._progress_async_consensus() + + assert not transceiver.cancel_request(req) + assert coordinator.ready_withdrawals == [] + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_RELEASE, + 34, + 0, + ConsensusOutcome.READY, + ) + ) + transceiver._progress_async_consensus() + assert not transceiver.cancel_request(req) + + transceiver.activate_context_requests_for_schedule([req]) + assert not transceiver.cancel_request(req) + + coordinator.events.append( + ConsensusEvent( + ConsensusEventKind.READY_COMPLETE, + 34, + 0, + ConsensusOutcome.READY, + ) + ) + transceiver._progress_async_consensus() + assert transceiver.cancel_request(req) diff --git a/tests/unittest/disaggregated/test_transfer_quiescence.py b/tests/unittest/disaggregated/test_transfer_quiescence.py new file mode 100644 index 000000000000..8be6ecf5ef67 --- /dev/null +++ b/tests/unittest/disaggregated/test_transfer_quiescence.py @@ -0,0 +1,1442 @@ +# 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. + +import threading +import weakref +from collections import OrderedDict +from types import SimpleNamespace + +import numpy as np +import pytest + +from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus +from tensorrt_llm._torch.disaggregation.native import transfer as transfer_module +from tensorrt_llm._torch.disaggregation.native.transfer import ( + AgentResult, + KVSendTask, + MessageType, + Receiver, + RecvReqInfo, + RxSession, + Sender, + TaskStatus, + TransferWorker, + TxSession, + _ControlPlane, + _decode_protocol_capabilities, + _encode_protocol_capabilities, +) + + +class _Task: + def __init__(self, status: TaskStatus, *, expected_transfers: int = 1, slice_id: int = 0): + self.status = status + self.expected_transfers = expected_transfers + self.slice_id = slice_id + self.last_slice_count = 0 + self._perf_timer = None + self._event = threading.Event() + if status in (TaskStatus.TRANSFERRED, TaskStatus.ERROR): + self._event.set() + + @property + def is_done(self) -> bool: + return self._event.is_set() + + def fail(self, _error: Exception) -> None: + self.status = TaskStatus.ERROR + self._event.set() + + def complete(self) -> None: + self.status = TaskStatus.TRANSFERRED + self._event.set() + + def print_perf_info(self, *_args) -> None: + pass + + +class _Sender: + def __init__(self): + self.cancelled: list[int] = [] + self.acks: list[tuple[str, int, bool]] = [] + self.ack_epochs: list[int | None] = [] + self.cleared: list[int] = [] + + def send_cancel_to_receivers(self, unique_rid: int) -> None: + self.cancelled.append(unique_rid) + + def send_cancel_ack( + self, + endpoint: str, + unique_rid: int, + *, + request_epoch=None, + from_worker: bool, + ) -> bool: + self.acks.append((endpoint, unique_rid, from_worker)) + self.ack_epochs.append(request_epoch) + return True + + def clear_session(self, unique_rid: int, *_args) -> None: + self.cleared.append(unique_rid) + + def _get_req_info(self, _unique_rid: int) -> dict: + return {} + + def dispatch_task(self, _task, _snapshot, *, operation_owner: TxSession) -> None: + operation_owner.finish_dispatch([]) + + +class _Bounce: + def __init__(self): + self.released: list[tuple[int, int]] = [] + self.orphaned: list[tuple[int, int]] = [] + self.drained: list[tuple[int, int]] = [] + self.settlement_callbacks = {} + + def release_idle_reservation(self, rid_slice: tuple[int, int]) -> None: + self.released.append(rid_slice) + + def orphan_reservation(self, rid_slice: tuple[int, int]) -> None: + if rid_slice not in self.orphaned: + self.orphaned.append(rid_slice) + + def confirm_drained(self, rid_slice: tuple[int, int]) -> None: + self.drained.append(rid_slice) + callback = self.settlement_callbacks.pop(rid_slice, None) + if callback is not None: + callback(False) + + def set_completion_callback(self, rid_slice, on_settled) -> None: + assert rid_slice not in self.settlement_callbacks + self.settlement_callbacks[rid_slice] = on_settled + + def is_bounced(self, _rid_slice: tuple[int, int]) -> bool: + return False + + def record_result(self, *_args, **_kwargs) -> None: + callback = _kwargs.get("on_done") + if callback is None and len(_args) >= 6: + callback = _args[5] + if callback is not None: + callback(True) + + def record_failure(self, *_args, **_kwargs) -> None: + callback = _kwargs.get("on_done") + if callback is not None: + callback(False) + + +class _InlineBounce(_Bounce): + def is_bounced(self, _rid_slice: tuple[int, int]) -> bool: + return True + + +class _RaisingBounce(_InlineBounce): + def record_result(self, *_args, **_kwargs) -> None: + raise RuntimeError("injected scatter setup failure") + + +class _Receiver: + endpoint = "tcp://receiver" + + def __init__(self): + self._bounce = _Bounce() + self._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(instance_name="gen", instance_rank=0) + ) + self.cancelled: list[tuple[int, set[str], set[str]]] = [] + self.retained: list[RxSession] = [] + self.cleared: list[int] = [] + + def send_cancel_to_senders( + self, + unique_rid: int, + endpoints: set[str], + ack_capable_endpoints: set[str], + ) -> None: + self.cancelled.append((unique_rid, endpoints, ack_capable_endpoints)) + + def retain_draining_session(self, session: RxSession) -> None: + self.retained.append(session) + + def clear_session(self, unique_rid: int, *_args) -> None: + self.cleared.append(unique_rid) + self.retained.clear() + + def _fail_protocol(self, _error: RuntimeError) -> None: + pass + + +class _AuxBuffer: + def __init__(self): + self.freed: list[int] = [] + + def free_slot(self, slot: int) -> None: + self.freed.append(slot) + + +class _ShutdownMessenger: + endpoint = "tcp://receiver" + + def __init__(self, order: list[str]): + self._listener_thread = None + self._order = order + self.stopped = threading.Event() + + def stop(self) -> None: + self._order.append("listener-stop") + self.stopped.set() + + +class _ShutdownBounce(_Bounce): + def __init__(self, order: list[str]): + super().__init__() + self._order = order + self._scatter_thread = None + self.closed = threading.Event() + + def close(self) -> None: + self._order.append("bounce-close") + self.closed.set() + + +class _ShutdownAgent: + def __init__(self, order: list[str]): + self._order = order + + def deregister_memory(self, _desc) -> None: + self._order.append("deregister-memory") + + def shutdown(self) -> None: + self._order.append("agent-shutdown") + + +def _make_tx_session(task_status: TaskStatus = TaskStatus.INIT) -> tuple[TxSession, object]: + task = _Task(task_status) + session = object.__new__(TxSession) + session.lock = threading.Lock() + session._sealed = False + session._closed = False + session._close_requested = False + session._terminal_status = None + session._terminal_snapshot = None + session._exception = None + session._outstanding_operations = 0 + session._cancel_ack_endpoints = set() + session._cancel_acked_endpoints = set() + session._request_operations = set() + session._cancel_ack_operations = set() + session._cancel_acked_operations = set() + session._cancel_notified = False + session._need_aux = False + session.kv_tasks = [task] + session.aux_task = None + session._sender = _Sender() + session._aux_buffer = None + session.aux_slot = None + session._base_args = SimpleNamespace( + params=SimpleNamespace(disagg_request_id=11, ctx_request_id=None), + prompt_len=None, + beam_width=1, + ) + session.request_id = 11 + return session, task + + +def _make_rx_session(task_status: TaskStatus = TaskStatus.INIT) -> tuple[RxSession, object]: + task = _Task(task_status) + session = object.__new__(RxSession) + session.lock = threading.Lock() + session._sealed = False + session._closed = False + session._close_requested = False + session._terminal_status = None + session._exception = None + session.transfer_start_time = None + session.transfer_end_time = None + session.kv_cache_size_bytes = 0 + session._kv_tasks = [task] + session._need_aux = False + session._outstanding_operations = 0 + session._expected_operations = {} + session._retired_operation_keys = set() + session._pending_scatter_callbacks = 0 + session._pending_bounce_settlements = set() + session._aux_obligations_reserved = False + session._aux_obligation_slice_id = None + session._aux_count = 0 + session._aux_status = TaskStatus.INIT + session._sender_endpoints = set() + session._cancel_pending_endpoints = set() + session._cancel_pending_operations = set() + session._ack_capable_endpoints = set() + session._v2_enabled = None + session._legacy_rank_cohorts = {} + session._legacy_bound_cohorts = {} + session.request_epoch = 123 + session._receiver = _Receiver() + session._aux_buffer = None + session.aux_slot = None + session._base_args = SimpleNamespace( + params=SimpleNamespace(disagg_request_id=17, ctx_request_id=None) + ) + session.request_id = 17 + return session, task + + +def test_tx_seal_rejects_an_already_credited_queued_operation() -> None: + session, task = _make_tx_session() + session._outstanding_operations = 1 + + assert not session.seal_and_check_quiescent() + assert not session.try_mark_transferring(task) + assert task.status == TaskStatus.INIT + + session.retire_operation() + assert not session.has_transferring_tasks() + + +def test_tx_quiescence_includes_aux_transfer() -> None: + session, _ = _make_tx_session(TaskStatus.TRANSFERRED) + session.aux_task = SimpleNamespace(status=TaskStatus.TRANSFERRING) + session._outstanding_operations = 1 + + assert not session.seal_and_check_quiescent() + assert session.has_transferring_tasks() + + +def test_tx_exception_does_not_hide_an_active_transfer() -> None: + session, task = _make_tx_session(TaskStatus.TRANSFERRING) + session._outstanding_operations = 2 + + session.set_exception("peer transfer failed") + + assert task.status == TaskStatus.TRANSFERRING + assert not session.seal_and_check_quiescent() + + session.retire_operation() + assert session.has_transferring_tasks() + session.retire_operation() + assert not session.has_transferring_tasks() + + +def test_tx_cancel_ack_waits_for_every_peer_obligation() -> None: + session, _ = _make_tx_session(TaskStatus.TRANSFERRING) + session._outstanding_operations = 2 + + session.cancel(ack_endpoint="tcp://receiver") + session.retire_operation() + assert session._sender.acks == [] + + session.retire_operation() + assert session._sender.acks == [("tcp://receiver", 11, True)] + + session.cancel(ack_endpoint="tcp://receiver") + assert session._sender.acks == [ + ("tcp://receiver", 11, True), + ("tcp://receiver", 11, False), + ] + assert session._sender.cancelled == [11] + + +def test_tx_v2_cancel_after_terminal_snapshot_still_acknowledges_drain() -> None: + session, _ = _make_tx_session(TaskStatus.TRANSFERRED) + session.register_request_operation("tcp://receiver", 123) + session._terminal_snapshot = SessionStatus.KV_TRANSFERRED + + session.cancel(ack_endpoint="tcp://receiver", request_epoch=123) + + assert session._sender.acks == [("tcp://receiver", 11, False)] + assert session._sender.ack_epochs == [123] + + # The receiver retransmits CANCEL until it observes an ACK. A duplicate + # request after drain must resend the same idempotent epoch-bound ACK. + session.cancel(ack_endpoint="tcp://receiver", request_epoch=123) + assert session._sender.acks == [ + ("tcp://receiver", 11, False), + ("tcp://receiver", 11, False), + ] + assert session._sender.ack_epochs == [123, 123] + + +def test_tx_v2_cancel_tracks_every_receiver_operation() -> None: + session, _ = _make_tx_session(TaskStatus.TRANSFERRED) + session.register_request_operation("tcp://receiver-0", 123) + session.register_request_operation("tcp://receiver-1", 123) + + session.cancel(ack_endpoint="tcp://receiver-0", request_epoch=123) + session.cancel(ack_endpoint="tcp://receiver-1", request_epoch=123) + + assert set(session._sender.acks) == { + ("tcp://receiver-0", 11, False), + ("tcp://receiver-1", 11, False), + } + assert session._sender.ack_epochs == [123, 123] + + +def test_tx_close_defers_aux_reclamation_until_peer_response() -> None: + session, _ = _make_tx_session(TaskStatus.TRANSFERRING) + aux_buffer = _AuxBuffer() + session._aux_buffer = aux_buffer + session.aux_slot = 7 + session._outstanding_operations = 1 + + session.close() + assert aux_buffer.freed == [] + + session.retire_operation() + assert aux_buffer.freed == [7] + assert session._sender.cleared == [11] + + +def test_tx_terminal_snapshot_does_not_freeze_pending_dispatch() -> None: + session, task = _make_tx_session(TaskStatus.TRANSFERRED) + session._outstanding_operations = 1 + + assert session.seal_and_snapshot_terminal() is None + assert not session._sealed + + session.retire_operation() + assert session.seal_and_snapshot_terminal() == SessionStatus.KV_TRANSFERRED + assert task.status == TaskStatus.TRANSFERRED + + +def test_tx_terminal_snapshot_is_immutable_across_cancel_race() -> None: + for _ in range(100): + session, _ = _make_tx_session(TaskStatus.TRANSFERRED) + barrier = threading.Barrier(3) + + def snapshot() -> None: + barrier.wait() + session.seal_and_snapshot_terminal() + + def cancel() -> None: + barrier.wait() + session.cancel() + + snapshot_thread = threading.Thread(target=snapshot) + cancel_thread = threading.Thread(target=cancel) + snapshot_thread.start() + cancel_thread.start() + barrier.wait() + snapshot_thread.join() + cancel_thread.join() + + snapshot_status = session.seal_and_snapshot_terminal() + assert snapshot_status in (SessionStatus.KV_TRANSFERRED, SessionStatus.CANCELLED) + assert session.status == snapshot_status + + +def test_tx_pre_cancel_rejects_kv_and_aux_without_throwing() -> None: + session, _ = _make_tx_session() + session.kv_tasks.clear() + + session.cancel() + session.send(object()) + aux_task = session.send_aux() + + assert session.kv_tasks[-1].status == TaskStatus.ERROR + assert aux_task.status == TaskStatus.ERROR + + +def test_tx_cancel_between_kv_and_aux_keeps_both_terminal() -> None: + session, _ = _make_tx_session() + session.kv_tasks.clear() + + session.send(object()) + session.cancel() + aux_task = session.send_aux() + + assert session.kv_tasks[0].status == TaskStatus.ERROR + assert aux_task.status == TaskStatus.ERROR + + +def test_rx_seal_blocks_transfer_start_and_endpoint_publication() -> None: + session, task = _make_rx_session() + + assert session.seal_and_check_quiescent() + assert not session.mark_transferring(0, {"tcp://sender"}) + assert task.status == TaskStatus.INIT + assert session._sender_endpoints == set() + + +def test_rx_seal_waits_for_started_transfer() -> None: + session, task = _make_rx_session() + task.expected_transfers = 1 + + assert session.mark_transferring( + 0, + {0: "tcp://sender"}, + {"tcp://sender"}, + request_epoch=session.request_epoch, + ) + assert not session.seal_and_check_quiescent() + assert task.status == TaskStatus.TRANSFERRING + assert session._sender_endpoints == {"tcp://sender"} + + +@pytest.mark.parametrize( + "peer_rank,sender_endpoint,request_epoch,error", + [ + (1, "tcp://sender", 123, "unexpected native-transfer result operation"), + (0, "tcp://other", 123, "result source mismatch"), + (0, "tcp://sender", 124, "result epoch mismatch"), + ], +) +def test_rx_wrong_v2_result_identity_fails_closed_without_retiring_credit( + peer_rank, sender_endpoint, request_epoch, error +) -> None: + session, task = _make_rx_session() + task.expected_transfers = 1 + assert session.mark_transferring( + 0, + {0: "tcp://sender"}, + {"tcp://sender"}, + request_epoch=session.request_epoch, + ) + + session.process_kv_agent_result( + peer_rank, + 0, + True, + AgentResult.SUCCESS, + transfer_size=4096, + sender_endpoint=sender_endpoint, + request_epoch=request_epoch, + ) + + assert session.kv_cache_size_bytes == 0 + assert session._outstanding_operations == 1 + assert task.status == TaskStatus.TRANSFERRING + assert session.status == SessionStatus.ERROR + assert error in str(session.exception) + assert session._receiver.cancelled == [(17, {"tcp://sender"}, {("tcp://sender", 123)})] + + +def test_rx_exact_v2_result_retires_once_and_duplicate_is_idempotent() -> None: + session, task = _make_rx_session() + task.expected_transfers = 1 + assert session.mark_transferring( + 0, + {0: "tcp://sender"}, + {"tcp://sender"}, + request_epoch=session.request_epoch, + ) + + for _ in range(2): + session.process_kv_agent_result( + 0, + 0, + True, + AgentResult.SUCCESS, + transfer_size=4096, + sender_endpoint="tcp://sender", + request_epoch=session.request_epoch, + ) + + assert session.kv_cache_size_bytes == 4096 + assert session._outstanding_operations == 0 + assert task.status == TaskStatus.TRANSFERRED + assert session.transfer_end_time is not None + assert session._retired_operation_keys == {("kv", 0, 0)} + + +def test_rx_generation_first_aux_obligations_block_quiescence() -> None: + session, task = _make_rx_session() + session._need_aux = True + task.expected_transfers = 2 + + assert session.mark_transferring(0, {"tcp://sender-0", "tcp://sender-1"}) + assert session._outstanding_operations == 4 + with session.lock: + session._retire_receive_operation_unlocked(("kv", 0, 0)) + session._retire_receive_operation_unlocked(("kv", 0, 1)) + + assert session._outstanding_operations == 2 + assert session.has_transferring_tasks() + + +def test_rx_cancel_ack_defers_aux_and_bounce_reclamation_until_drain() -> None: + session, task = _make_rx_session() + session._need_aux = True + task.expected_transfers = 1 + aux_buffer = _AuxBuffer() + session._aux_buffer = aux_buffer + session.aux_slot = 3 + + assert session.mark_transferring( + 0, + {0: "tcp://sender"}, + {"tcp://sender"}, + request_epoch=session.request_epoch, + ) + session.cancel() + session.close() + + assert aux_buffer.freed == [] + assert session._receiver.cleared == [] + assert session._receiver._bounce.orphaned == [(17, 0)] + + session.process_cancel_ack("tcp://sender", session.request_epoch) + + assert aux_buffer.freed == [3] + assert session._receiver.cleared == [17] + assert session._receiver._bounce.drained == [] + + +def test_rx_bounced_cancel_before_result_retires_unconditional_settlement_credit() -> None: + session, task = _make_rx_session() + task.expected_transfers = 1 + assert session.mark_transferring( + 0, + {0: "tcp://sender"}, + {"tcp://sender"}, + request_epoch=session.request_epoch, + bounced=True, + ) + rid_slice = (session.disagg_request_id, task.slice_id) + session._receiver._bounce.set_completion_callback( + rid_slice, session._make_bounce_settlement_callback(task) + ) + + session.cancel() + session.close() + assert session._pending_bounce_settlements == {0} + assert session._receiver.cleared == [] + + session.process_cancel_ack("tcp://sender", session.request_epoch) + + assert session._receiver._bounce.drained == [rid_slice] + assert session._pending_bounce_settlements == set() + assert session._receiver.cleared == [17] + + +def test_rx_malformed_bounce_result_keeps_credit_until_exact_ack_settles() -> None: + session, task = _make_rx_session() + task.expected_transfers = 1 + session._receiver._bounce = _RaisingBounce() + session._receiver._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(instance_name="gen", instance_rank=0) + ) + assert session.mark_transferring( + 0, + {0: "tcp://sender"}, + {"tcp://sender"}, + request_epoch=session.request_epoch, + bounced=True, + ) + rid_slice = (session.disagg_request_id, task.slice_id) + session._receiver._bounce.set_completion_callback( + rid_slice, session._make_bounce_settlement_callback(task) + ) + + session.process_kv_agent_result( + 0, + 0, + True, + AgentResult.SUCCESS, + sender_endpoint="tcp://sender", + request_epoch=session.request_epoch, + ) + + assert task.status == TaskStatus.ERROR + assert session._outstanding_operations == 1 + assert session._pending_bounce_settlements == {0} + session.close() + + session.process_cancel_ack("tcp://sender", session.request_epoch) + + assert session._receiver._bounce.drained == [rid_slice] + assert session._pending_bounce_settlements == set() + assert session._receiver.cleared == [17] + + +def test_rx_legacy_cancel_finalizes_after_kv_and_aux_terminal_results() -> None: + session, task = _make_rx_session() + session._need_aux = True + task.expected_transfers = 1 + aux_buffer = _AuxBuffer() + session._aux_buffer = aux_buffer + session.aux_slot = 4 + + assert session.mark_transferring(0, {"tcp://legacy-sender"}, set()) + session.cancel() + session.close() + + session.process_kv_agent_result(0, 0, True, AgentResult.FAILED) + assert aux_buffer.freed == [] + session.process_aux_agent_result(0, AgentResult.FAILED) + + assert aux_buffer.freed == [4] + assert session._receiver.cleared == [17] + + +def test_rx_bounce_fallback_inline_callback_does_not_reenter_session_lock() -> None: + session, task = _make_rx_session(TaskStatus.TRANSFERRING) + task.expected_transfers = 1 + session._outstanding_operations = 1 + session._receiver._bounce = _InlineBounce() + session._receiver._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(instance_name="gen", instance_rank=0) + ) + + finished = threading.Event() + + def deliver_result() -> None: + session.process_kv_agent_result( + 0, + 0, + True, + AgentResult.SUCCESS, + ) + finished.set() + + thread = threading.Thread(target=deliver_result) + thread.start() + thread.join(timeout=1) + + assert finished.is_set(), "inline bounce completion deadlocked on RxSession.lock" + assert task.status == TaskStatus.TRANSFERRED + assert session._outstanding_operations == 0 + assert session._pending_scatter_callbacks == 0 + + +def test_rx_scatter_setup_failure_releases_only_after_terminal_result_proof() -> None: + session, task = _make_rx_session(TaskStatus.TRANSFERRING) + task.expected_transfers = 1 + session._outstanding_operations = 1 + session._receiver._bounce = _RaisingBounce() + session._pending_bounce_settlements = {0} + session._pending_scatter_callbacks = 1 + session._receiver._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(instance_name="gen", instance_rank=0) + ) + + session.process_kv_agent_result(0, 0, True, AgentResult.SUCCESS) + + assert task.status == TaskStatus.ERROR + assert session._outstanding_operations == 1 + assert session._pending_scatter_callbacks == 1 + assert session._receiver._bounce.orphaned == [(17, 0)] + assert session._receiver._bounce.drained == [] + + +def test_rx_scatter_callback_credit_defers_close() -> None: + session, task = _make_rx_session(TaskStatus.TRANSFERRED) + task.expected_transfers = 1 + session._pending_scatter_callbacks = 1 + aux_buffer = _AuxBuffer() + session._aux_buffer = aux_buffer + session.aux_slot = 5 + + session.close() + assert aux_buffer.freed == [] + + session._retire_scatter_callback() + assert aux_buffer.freed == [5] + + +def test_worker_shutdown_keeps_rx_progress_and_memory_alive_until_cancel_ack() -> None: + order: list[str] = [] + bounce = _ShutdownBounce(order) + messenger = _ShutdownMessenger(order) + + receiver = object.__new__(Receiver) + receiver._shutdown = False + receiver._shutdown_lock = threading.Lock() + receiver._shutdown_complete = threading.Event() + receiver._shutdown_thread = None + receiver._shutdown_error = None + receiver._messenger = messenger + receiver._bounce = bounce + receiver._control = SimpleNamespace(flush=lambda: None, shutdown=lambda: None) + receiver._sessions = {} + receiver._sessions_lock = threading.Lock() + receiver._sessions_drained = threading.Condition(receiver._sessions_lock) + receiver._draining_sessions = {} + receiver._pre_cancelled_rids = {} + receiver._closed_rids = OrderedDict() + + session, task = _make_rx_session(TaskStatus.TRANSFERRING) + task.expected_transfers = 1 + session._receiver = receiver + session._terminal_status = SessionStatus.CANCELLED + session._outstanding_operations = 1 + session._sender_endpoints = {"tcp://sender"} + session._ack_capable_endpoints = {"tcp://sender"} + session._cancel_pending_endpoints = {"tcp://sender"} + aux_buffer = _AuxBuffer() + session._aux_buffer = aux_buffer + session.aux_slot = 9 + receiver._sessions[session.disagg_request_id] = weakref.ref(session) + + drain_started = threading.Event() + retain_draining_session = receiver.retain_draining_session + + def retain_and_signal(rx_session: RxSession) -> None: + retain_draining_session(rx_session) + drain_started.set() + + receiver.retain_draining_session = retain_and_signal + + worker = object.__new__(TransferWorker) + worker._shutdown = False + worker._shutdown_lock = threading.Lock() + worker._shutdown_complete = threading.Event() + worker._shutdown_thread = None + worker._shutdown_error = None + worker._rank_info_server = None + worker._sender = SimpleNamespace( + shutdown=lambda: order.append("sender-stop"), + _worker_threads=[], + _messenger=None, + ) + worker._receiver = receiver + worker._bounce = bounce + worker._registered_mem = [object()] + worker._agent = _ShutdownAgent(order) + + # Model shutdown requested from a bounce completion callback. It must + # return and let that callback process the final ACK rather than joining + # itself or blocking the callback's operation-credit retirement. + bounce._scatter_thread = threading.current_thread() + shutdown_complete = worker.shutdown() + + assert drain_started.wait(timeout=1) + assert shutdown_complete is worker._shutdown_complete + assert not messenger.stopped.is_set() + assert not bounce.closed.is_set() + assert aux_buffer.freed == [] + assert "deregister-memory" not in order + + session.process_cancel_ack("tcp://sender") + + assert worker._shutdown_complete.wait(timeout=1) + assert session._closed + assert aux_buffer.freed == [9] + assert order == [ + "sender-stop", + "listener-stop", + "bounce-close", + "deregister-memory", + "agent-shutdown", + ] + + +def test_rx_pre_cancel_receive_is_nonthrowing() -> None: + session, _ = _make_rx_session() + session._kv_tasks.clear() + + session.cancel() + session.receive(object()) + + assert session._kv_tasks[-1].status == TaskStatus.ERROR + + +def test_sender_closed_tombstone_rejects_late_request() -> None: + sender = object.__new__(Sender) + sender._shutdown = True + sender._sessions = {} + sender._sessions_lock = threading.Lock() + sender._closed_rids = OrderedDict([(29, None)]) + sender._pre_cancelled_rids = {} + sender._peer_capabilities = {} + sender._registrar = SimpleNamespace( + get_peer_rank_info=lambda _name, _rank: SimpleNamespace(self_endpoint="tcp://receiver") + ) + failed: list[int] = [] + saved: list[int] = [] + sender._send_failed_result_to_receiver = lambda info, **_kwargs: failed.append(info.unique_rid) + sender._save_peer_req_info = lambda info: saved.append(info.unique_rid) + info = RecvReqInfo( + sender_req_id=1, + instance_name="gen", + instance_rank=0, + block_ids_per_layer_groups=[np.array([], dtype=np.int64)], + unique_rid=29, + ) + + sender._respond_with_kv(b"", [MessageType.REQUEST_DATA, info.to_bytes()]) + + assert failed == [29] + assert saved == [] + + +def test_closed_tombstones_are_bounded() -> None: + tombstones: OrderedDict[int, None] = OrderedDict() + original_limit = Receiver._TOMBSTONE_LIMIT + Receiver._TOMBSTONE_LIMIT = 3 + try: + for request_id in range(5): + Receiver._remember_tombstone(tombstones, request_id) + finally: + Receiver._TOMBSTONE_LIMIT = original_limit + + assert list(tombstones) == [2, 3, 4] + + +def test_receiver_epoch_tombstone_drops_only_matching_delayed_result() -> None: + receiver = object.__new__(Receiver) + receiver._sessions = {} + receiver._sessions_lock = threading.Lock() + receiver._closed_operations = OrderedDict([((17, 123), None)]) + receiver._protocol_error = None + matching = transfer_module._make_kv_result_msg( + 0, + 17, + 0, + True, + AgentResult.SUCCESS, + request_epoch=123, + sender_endpoint="tcp://sender", + ) + receiver._process_kv_agent_result(b"", matching) + assert receiver._protocol_error is None + + stale = transfer_module._make_kv_result_msg( + 0, + 17, + 0, + True, + AgentResult.SUCCESS, + request_epoch=124, + sender_endpoint="tcp://sender", + ) + receiver._process_kv_agent_result(b"", stale) + assert receiver._protocol_error is not None + assert "unknown request incarnation" in str(receiver._protocol_error) + + +def test_sender_acknowledged_pre_cancel_is_not_count_evicted() -> None: + sender = object.__new__(Sender) + sender._sessions = {} + sender._sessions_lock = threading.Lock() + sender._ingress_lock = threading.Lock() + sender._shutdown = False + sender._pre_cancelled_rids = {} + sender._pre_cancelled_operations = {} + sender._cancelled_operation_tombstones = OrderedDict() + sender._closed_rids = OrderedDict() + sender._protocol_error = None + sender._peer_requests = {} + sender._peer_requests_lock = threading.Lock() + + original_limit = Sender._TOMBSTONE_LIMIT + Sender._TOMBSTONE_LIMIT = 3 + try: + for request_id in range(5): + sender._remember_pre_cancelled_unlocked(request_id) + finally: + Sender._TOMBSTONE_LIMIT = original_limit + + class _DelayedSession: + disagg_request_id = 0 + + def __init__(self): + self.cancelled = False + + def cancel(self) -> None: + self.cancelled = True + + session = _DelayedSession() + sender.setup_session(session) + + assert session.cancelled + assert 0 not in sender._pre_cancelled_rids + + +def test_sender_metadata_lease_survives_stale_sweep( + monkeypatch: pytest.MonkeyPatch, +) -> None: + sender = object.__new__(Sender) + sender._sessions = {} + sender._sessions_lock = threading.Lock() + sender._peer_requests = {1: {0: object()}, 2: {0: object()}} + sender._peer_requests_timestamps = {1: 1.0, 2: 1.0} + sender._peer_requests_lock = threading.Lock() + sender._peer_request_leases = {1} + monkeypatch.setattr(transfer_module.time, "monotonic", lambda: 1000.0) + + sender.sweep_stale_req_infos() + + assert 1 in sender._peer_requests + assert 1 in sender._peer_requests_timestamps + assert 2 not in sender._peer_requests + assert 2 not in sender._peer_requests_timestamps + + +def test_sender_stale_sweep_rechecks_lease_before_delete( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _GateLock: + def __init__(self) -> None: + self.entered = threading.Event() + self.release = threading.Event() + + def __enter__(self): + self.entered.set() + assert self.release.wait(timeout=1) + return self + + def __exit__(self, *_args) -> None: + pass + + sender = object.__new__(Sender) + sessions_lock = _GateLock() + sender._sessions = {} + sender._sessions_lock = sessions_lock + sender._peer_requests = {3: {0: object()}} + sender._peer_requests_timestamps = {3: 1.0} + sender._peer_requests_lock = threading.Lock() + sender._peer_request_leases = set() + sender._protocol_error = None + monkeypatch.setattr(transfer_module.time, "monotonic", lambda: 1000.0) + + sweep = threading.Thread(target=sender.sweep_stale_req_infos) + sweep.start() + assert sessions_lock.entered.wait(timeout=1) + sender.pin_peer_req_infos(3) + sessions_lock.release.set() + sweep.join(timeout=1) + + assert not sweep.is_alive() + assert 3 in sender._peer_requests + assert 3 in sender._peer_request_leases + + +def test_sender_metadata_only_cancel_tombstones_and_notifies() -> None: + sender = object.__new__(Sender) + sender._ingress_lock = threading.Lock() + sender._sessions = {} + sender._sessions_lock = threading.Lock() + metadata = object() + sender._peer_requests = {4: {0: metadata}} + sender._peer_requests_timestamps = {4: 1.0} + sender._peer_requests_lock = threading.Lock() + sender._peer_request_leases = {4} + sender._pre_cancelled_rids = {} + sender._metadata_cancelled_rids = OrderedDict() + sender._protocol_error = None + sender._closed_rids = OrderedDict() + notified = [] + sender._send_cancel_to_req_infos = lambda request_id, req_infos: notified.append( + (request_id, req_infos) + ) + + assert sender.cancel_peer_req_infos(4) + + assert list(sender._closed_rids) == [4] + assert 4 not in sender._pre_cancelled_rids + assert list(sender._metadata_cancelled_rids) == [4] + assert 4 not in sender._peer_requests + assert 4 not in sender._peer_requests_timestamps + assert 4 not in sender._peer_request_leases + assert notified == [(4, [metadata])] + + +def test_sender_metadata_only_cancel_history_is_bounded_and_cancels_delayed_session() -> None: + sender = object.__new__(Sender) + sender._sessions = {} + sender._sessions_lock = threading.Lock() + sender._ingress_lock = threading.Lock() + sender._peer_requests = {} + sender._peer_requests_timestamps = {} + sender._peer_requests_lock = threading.Lock() + sender._peer_request_leases = set() + sender._pre_cancelled_rids = {} + sender._metadata_cancelled_rids = OrderedDict() + sender._closed_rids = OrderedDict() + sender._protocol_error = None + sender._send_cancel_to_req_infos = lambda _request_id, _req_infos: None + + original_limit = Sender._TOMBSTONE_LIMIT + Sender._TOMBSTONE_LIMIT = 3 + try: + for request_id in range(5): + assert sender.cancel_peer_req_infos(request_id) + finally: + Sender._TOMBSTONE_LIMIT = original_limit + + assert not sender._pre_cancelled_rids + assert list(sender._metadata_cancelled_rids) == [2, 3, 4] + assert sender._protocol_error is None + + class _DelayedSession: + disagg_request_id = 4 + + def __init__(self): + self.cancelled = False + + def cancel(self) -> None: + self.cancelled = True + + session = _DelayedSession() + sender._shutdown = False + sender.setup_session(session) + + assert session.cancelled + assert 4 not in sender._metadata_cancelled_rids + + +def test_sender_session_creation_assumes_metadata_lease() -> None: + class _Session: + disagg_request_id = 5 + + def cancel(self) -> None: + raise AssertionError("session must not be pre-cancelled") + + sender = object.__new__(Sender) + sender._ingress_lock = threading.Lock() + sender._shutdown = False + sender._protocol_error = None + sender._sessions = {} + sender._sessions_lock = threading.Lock() + sender._peer_requests = {} + sender._peer_requests_timestamps = {} + sender._peer_requests_lock = threading.Lock() + sender._peer_request_leases = {5} + sender._pre_cancelled_rids = {} + sender._pre_cancelled_operations = {} + sender._cancelled_operation_tombstones = OrderedDict() + sender._closed_rids = OrderedDict() + session = _Session() + + sender.setup_session(session) + + assert sender._sessions[5]() is session + assert 5 not in sender._peer_request_leases + + +def test_sender_shutdown_gate_rejects_enqueue_after_worker_sentinels() -> None: + sender = object.__new__(Sender) + sender._shutdown = True + sender._ingress_lock = threading.Lock() + sender._num_threads = 1 + sender._send_task_queues = [SimpleNamespace(put=lambda _item: None)] + write_meta = SimpleNamespace(unique_rid=31, peer_rank=0) + + with pytest.raises(RuntimeError, match="shutting down"): + sender._enqueue(write_meta) + + +def test_sender_metadata_failure_retains_owner_if_notification_cannot_queue() -> None: + sender = object.__new__(Sender) + sender._stalled_operations_lock = threading.Lock() + sender._stalled_session_owners = [] + + def fail_metadata(_task, _info): + raise RuntimeError("injected metadata failure") + + sender._build_kv_write_meta = fail_metadata + sender._send_failed_result_to_receiver = lambda _info, **_kwargs: False + + task = object.__new__(KVSendTask) + task.status = TaskStatus.INIT + task._event = threading.Event() + task._exception = None + task._unique_rid = 11 + task._perf_timer = None + owner, _ = _make_tx_session() + owner._outstanding_operations = 1 + infos = { + 0: SimpleNamespace(instance_rank=0), + 1: SimpleNamespace(instance_rank=1), + } + + sender.dispatch_task(task, infos, operation_owner=owner) + + assert task.status == TaskStatus.ERROR + assert owner._outstanding_operations == 2 + assert sender._stalled_session_owners == [owner] + + +def test_native_protocol_capabilities_are_optional_and_backward_compatible() -> None: + legacy = _decode_protocol_capabilities(None) + negotiated = _decode_protocol_capabilities(_encode_protocol_capabilities()) + + assert legacy.version == 1 + assert not legacy.drain_ack + assert negotiated.version >= 2 + assert negotiated.drain_ack + + with pytest.raises(RuntimeError, match="invalid native-transfer capability"): + _decode_protocol_capabilities(b"not-a-capability-frame") + + +def test_explicit_async_mode_rejects_legacy_peer_before_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _LegacyInfoMessenger: + def __init__(self, *, mode: str, endpoint: str): + assert mode == "DEALER" + assert endpoint == "tcp://legacy-info" + + def send(self, message: list[bytes]) -> None: + assert message == [MessageType.REQUEST_INSTANCE_INFO] + + def receive(self) -> list[bytes]: + # No optional capability frame means protocol v1. + return [b"legacy-rank-info"] + + def stop(self) -> None: + pass + + receiver = object.__new__(Receiver) + receiver._sender_ep_instance_map = {} + receiver._sender_info_capabilities = {} + receiver._sender_endpoint_capabilities = {} + receiver._registrar = SimpleNamespace(self_rank_info=SimpleNamespace()) + sender_info = SimpleNamespace(sender_endpoints=["tcp://legacy-sender"]) + + monkeypatch.setattr(transfer_module, "ZMQMessenger", _LegacyInfoMessenger) + monkeypatch.setattr( + transfer_module, + "RankInfo", + SimpleNamespace(from_bytes=lambda _data: sender_info), + raising=False, + ) + monkeypatch.setenv("TRTLLM_PYTHON_TRANSCEIVER_ASYNC_CTX_TERMINAL_CONSENSUS", "1") + + with pytest.raises(RuntimeError, match="requires native-transfer protocol"): + receiver._get_sender_info(SimpleNamespace(ctx_info_endpoint="tcp://legacy-info")) + + +def test_control_plane_owns_dealers_on_one_thread_and_retries( + monkeypatch: pytest.MonkeyPatch, +) -> None: + created_on: list[int] = [] + sent_on: list[int] = [] + delivered: list[bytes] = [] + attempts: dict[bytes, int] = {} + + class _FakeMessenger: + def __init__(self, *, mode: str, endpoint: str): + assert mode == "DEALER" + assert endpoint == "tcp://peer" + created_on.append(threading.get_ident()) + + def send(self, message: list[bytes]) -> None: + sent_on.append(threading.get_ident()) + marker = message[0] + attempts[marker] = attempts.get(marker, 0) + 1 + if marker == b"retry" and attempts[marker] == 1: + raise RuntimeError("injected transient send failure") + delivered.append(marker) + + def stop(self) -> None: + assert threading.get_ident() in created_on + + monkeypatch.setattr(transfer_module, "ZMQMessenger", _FakeMessenger) + control = _ControlPlane("test") + callback_count = 0 + callback_lock = threading.Lock() + + def on_sent() -> None: + nonlocal callback_count + with callback_lock: + callback_count += 1 + + callers = [ + threading.Thread( + target=control.send, + args=("tcp://peer", [marker]), + kwargs={"retry": marker == b"retry", "on_sent": on_sent}, + ) + for marker in (b"retry", b"a", b"b", b"c") + ] + for caller in callers: + caller.start() + for caller in callers: + caller.join(timeout=1) + + assert all(not caller.is_alive() for caller in callers) + control.flush() + owner_ident = control.owner_ident + control.shutdown() + + assert attempts[b"retry"] == 2 + assert sorted(delivered) == [b"a", b"b", b"c", b"retry"] + assert callback_count == 4 + assert owner_ident is not None + assert set(created_on) == {owner_ident} + assert set(sent_on) == {owner_ident} + + +def test_control_plane_retransmits_cancel_until_ack_predicate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + attempts = 0 + acknowledged = threading.Event() + + class _FakeMessenger: + def __init__(self, *, mode: str, endpoint: str): + assert mode == "DEALER" + assert endpoint == "tcp://peer" + + def send(self, _message: list[bytes]) -> None: + nonlocal attempts + attempts += 1 + + def stop(self) -> None: + pass + + monkeypatch.setattr(transfer_module, "ZMQMessenger", _FakeMessenger) + control = _ControlPlane("cancel-retransmit") + control.send( + "tcp://peer", + [MessageType.CANCEL_SESSION], + retry=True, + wait=False, + repeat_until=acknowledged.is_set, + ) + deadline = threading.Event() + for _ in range(100): + if attempts >= 2: + break + deadline.wait(0.005) + assert attempts >= 2 + + acknowledged.set() + control.flush() + settled_attempts = attempts + deadline.wait(0.02) + control.shutdown() + + assert attempts == settled_attempts + + +def test_rx_partial_fanout_waits_for_all_negotiated_drain_acks() -> None: + session, task = _make_rx_session() + task.expected_transfers = 2 + endpoints = {"tcp://sender-0", "tcp://sender-1"} + + assert session.mark_transferring( + 0, + {0: "tcp://sender-0", 1: "tcp://sender-1"}, + endpoints, + request_epoch=session.request_epoch, + ) + cancel_endpoints, cancel_operations, orphaned = session.fail_partial_dispatch( + 0, + {0}, + RuntimeError("injected fan-out failure"), + ) + + assert cancel_endpoints == {"tcp://sender-0"} + assert cancel_operations == {("tcp://sender-0", session.request_epoch)} + assert orphaned == [(17, 0)] + assert session._outstanding_operations == 1 + assert task.status == TaskStatus.ERROR + + session.process_cancel_ack("tcp://sender-0", session.request_epoch) + assert session._outstanding_operations == 0 + assert session._receiver._bounce.drained == [] + + +def test_rx_partial_later_slice_does_not_retire_first_slice_aux_credits() -> None: + session, first_task = _make_rx_session() + session._need_aux = True + first_task.expected_transfers = 2 + endpoints = {"tcp://sender-0", "tcp://sender-1"} + assert session.mark_transferring(0, endpoints, endpoints) + + second_task = _Task(TaskStatus.INIT, expected_transfers=2, slice_id=1) + session._kv_tasks.append(second_task) + assert session.mark_transferring(1, endpoints, endpoints) + assert session._outstanding_operations == 6 + + session.fail_partial_dispatch( + 1, + {"tcp://sender-0"}, + RuntimeError("injected second-slice fan-out failure"), + ) + + # Only one unsent KV writer from slice 1 is retired. Slice 0 owns both + # auxiliary credits and they remain live until result/ACK proof. + assert session._outstanding_operations == 5 + + +def test_rx_mixed_version_fanout_keeps_legacy_result_authoritative() -> None: + session, first_task = _make_rx_session() + first_task.expected_transfers = 1 + assert session.mark_transferring(0, {"tcp://new-sender"}, {"tcp://new-sender"}) + + second_task = _Task(TaskStatus.INIT, expected_transfers=1, slice_id=1) + session._kv_tasks.append(second_task) + assert session.mark_transferring(1, {"tcp://legacy-sender"}, set()) + + session.cancel() + + assert session._ack_capable_endpoints == set() + assert session._cancel_pending_endpoints == set() + assert session._receiver.cancelled == [ + ( + 17, + {"tcp://new-sender", "tcp://legacy-sender"}, + set(), + ) + ] + assert session._outstanding_operations == 2 + + +def test_rx_legacy_adp_cohort_is_bound_once_for_the_whole_request() -> None: + session, first_task = _make_rx_session() + cohorts = (frozenset({0, 1}), frozenset({2, 3})) + first_task.expected_transfers = 2 + assert session.mark_transferring( + 0, + {0: "tcp://sender-0", 1: "tcp://sender-1", 2: "tcp://sender-2", 3: "tcp://sender-3"}, + set(), + allowed_rank_cohorts=cohorts, + ) + with session.lock: + session._bind_legacy_cohort_unlocked(0, 0) + + second_task = _Task(TaskStatus.INIT, expected_transfers=2, slice_id=1) + session._kv_tasks.append(second_task) + assert session.mark_transferring( + 1, + {0: "tcp://sender-0", 1: "tcp://sender-1", 2: "tcp://sender-2", 3: "tcp://sender-3"}, + set(), + allowed_rank_cohorts=cohorts, + ) + with session.lock, pytest.raises(RuntimeError, match="different ADP writer cohort"): + session._bind_legacy_cohort_unlocked(1, 2) + + +def test_rx_legacy_adp_aux_first_binds_the_request_cohort() -> None: + session, task = _make_rx_session() + session._need_aux = True + cohorts = (frozenset({0, 1}), frozenset({2, 3})) + task.expected_transfers = 2 + assert session.mark_transferring( + 0, + {0: "tcp://sender-0", 1: "tcp://sender-1", 2: "tcp://sender-2", 3: "tcp://sender-3"}, + set(), + allowed_rank_cohorts=cohorts, + ) + + with session.lock: + assert ( + session._validate_receive_operation_unlocked( + ("aux", 2), sender_endpoint=None, request_epoch=None + ) + == "accept" + ) + with pytest.raises(RuntimeError, match="outside the bound ADP writer cohort"): + session._validate_receive_operation_unlocked( + ("kv", 0, 0), sender_endpoint=None, request_epoch=None + ) diff --git a/tests/unittest/llmapi/test_executor.py b/tests/unittest/llmapi/test_executor.py index dce923c6211c..4bf8d0804c00 100644 --- a/tests/unittest/llmapi/test_executor.py +++ b/tests/unittest/llmapi/test_executor.py @@ -18,6 +18,7 @@ GenerationRequest, GenerationResult, GenerationResultBase, PostprocWorker) from tensorrt_llm.executor.ipc import FusedIpcQueue, ZeroMqQueue +from tensorrt_llm.llmapi.llm import BaseLLM from tensorrt_llm.llmapi.tokenizer import TransformersTokenizer from tensorrt_llm.llmapi.utils import AsyncQueue from tensorrt_llm.sampling_params import SamplingParams @@ -29,6 +30,32 @@ WORLD_SIZE = mpi_world_size() +class _ChangingDisaggregatedParamsExecutor: + + def __init__(self) -> None: + self.params = {} + self.calls = 0 + + def get_disaggregated_params(self) -> dict: + self.calls += 1 + return self.params + + +def test_disaggregated_params_comes_from_current_executor_lifetime() -> None: + llm = BaseLLM.__new__(BaseLLM) + executor = _ChangingDisaggregatedParamsExecutor() + llm._executor = executor + + assert llm.disaggregated_params == {} + + executor.params = { + "ctx_info_endpoint": ["tcp://final:2000"], + "ctx_endpoint_generation": "final-generation", + } + assert llm.disaggregated_params == executor.params + assert executor.calls == 2 + + @pytest.fixture(scope="module") def engine_path(): return Path(tempfile.tempdir) / "llm_engine"