diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 6816e7bf9d7d..271f9a370d83 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -7,8 +7,9 @@ import threading import time from collections import defaultdict -from typing import Dict, List, Optional, Set, Tuple, Union +from typing import Dict, List, NamedTuple, Optional, Set, Tuple, Union +import msgspec import numpy as np import numpy.typing as npt import requests @@ -50,6 +51,9 @@ logger = logging.getLogger(__name__) +ABORT_TAG = b"ABORT" +ABORT_ACK_TAG = b"ABORT_ACK" + # Reuse a keep-alive session per bootstrap_addr for decode-side bootstrap queries # so we don't open a fresh TCP connection per query (that churns short-lived @@ -144,6 +148,116 @@ def __post_init__(self): self.rank_port = int(self.rank_port) +class AckTarget(NamedTuple): + ip: str + port: int + generation: int + + +class AbortNotification(msgspec.Struct, frozen=True): + room: int + decode_ip: Optional[str] = None + decode_port: Optional[int] = None + generation: Optional[int] = None + + @classmethod + def from_zmq(cls, msg: List[bytes]) -> Optional[AbortNotification]: + if len(msg) < 2: + logger.warning("Malformed ABORT message: too few frames (%d)", len(msg)) + return None + + try: + room = int(msg[1].decode("ascii")) + except (ValueError, UnicodeDecodeError): + logger.warning("Malformed ABORT message: invalid room field %r", msg[1]) + return None + + if len(msg) < 4: + return cls(room=room) + + try: + decode_ip = msg[2].decode("ascii") + decode_port = int(msg[3].decode("ascii")) + except (ValueError, UnicodeDecodeError): + logger.warning("Malformed ABORT message: invalid return address") + return cls(room=room) + + generation = None + if len(msg) >= 5: + try: + generation = int(msg[4].decode("ascii")) + except (ValueError, UnicodeDecodeError): + logger.warning("Malformed ABORT message: invalid generation") + return cls( + room=room, + decode_ip=decode_ip, + decode_port=decode_port, + generation=generation, + ) + + def to_zmq(self) -> List[bytes]: + frames = [ABORT_TAG, str(self.room).encode("ascii")] + if self.decode_ip is not None and self.decode_port is not None: + frames.extend( + [ + self.decode_ip.encode("ascii"), + str(self.decode_port).encode("ascii"), + ] + ) + if self.generation is not None: + frames.append(str(self.generation).encode("ascii")) + return frames + + def deferred_ack_target(self) -> Optional[AckTarget]: + if self.decode_ip is None or self.decode_port is None: + return None + if self.generation is None: + # Normal for a decode that did not arm (e.g. a prealloc abort before + # metadata was published): it does not wait for an ACK. An older + # decode that does wait falls back to its release timeout. + logger.debug( + "Generation-less ABORT for room %s; no deferred ACK", self.room + ) + return None + return AckTarget(self.decode_ip, self.decode_port, self.generation) + + +class AbortAck(msgspec.Struct, frozen=True): + room: int + prefill_rank: int + generation: Optional[int] = None + + @classmethod + def from_zmq(cls, msg: List[bytes]) -> Optional[AbortAck]: + if len(msg) < 3: + logger.warning("Incomplete ABORT_ACK received") + return None + try: + return cls( + room=int(msg[1].decode("ascii")), + prefill_rank=int(msg[2].decode("ascii")), + generation=(int(msg[3].decode("ascii")) if len(msg) >= 4 else None), + ) + except (ValueError, UnicodeDecodeError): + logger.warning("Malformed ABORT_ACK received") + return None + + def to_zmq(self) -> List[bytes]: + frames = [ + ABORT_ACK_TAG, + str(self.room).encode("ascii"), + str(self.prefill_rank).encode("ascii"), + ] + if self.generation is not None: + frames.append(str(self.generation).encode("ascii")) + return frames + + +class DeferredAbortAckState(msgspec.Struct): + generation: int + prefill_ranks: Set[int] + + class CommonKVManager(BaseKVManager): # Wire layout of the prefill->decode terminal status message. The legacy # layout (mooncake, and ascend which inherits it) is three untagged frames @@ -240,6 +354,12 @@ def __init__( self._socket_lock = threading.Lock() self.failure_records: Dict[int, str] = {} self.failure_lock = threading.Lock() + self._staging_outstanding: Dict[int, int] = defaultdict(int) + self._deferred_ack_targets: Dict[int, Dict[Tuple[str, int], AckTarget]] = {} + # A poisoned room has a permanently outstanding transfer whose backend + # can no longer prove quiescent. A new room lifecycle clears this state. + self._deferred_ack_poisoned_rooms: Set[int] = set() + if self.disaggregation_mode == DisaggregationMode.PREFILL: # When SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER is True, all CP ranks # participate in KV transfer; Otherwise only CP rank 0 sends. @@ -257,13 +377,6 @@ def __init__( ) self.register_to_bootstrap() self.transfer_infos = {} - # Deferred KV release: aborted room -> (decode_ip, decode_port); - # ack held until the transfer drains. - self._deferred_ack_targets: Dict[int, Tuple[str, int]] = {} - # Fan-out peers snapshotted at registration: the sender's clear() - # can pop transfer_infos before the worker drains, and the drain - # ack would otherwise reach the registered target alone. - self._deferred_ack_fanout_snapshots: Dict[int, List[Tuple[str, int]]] = {} self.req_to_decode_prefix_len: Dict[int, int] = {} self.decode_kv_args_table = {} self.pp_group = get_parallel().pp_group @@ -283,10 +396,11 @@ def __init__( self.session_pool_lock = threading.Lock() self.addr_to_rooms_tracker: Dict[str, Set[int]] = defaultdict(set) self.prefill_response_tracker: Dict[int, Set[int]] = defaultdict(set) - # Deferred KV release: room -> prefill ranks that acked their transfer - # drained. Entry exists only while the room is held, so a stale/late - # ack for a reused bootstrap_room is dropped. - self._deferred_abort_ack_tracker: Dict[int, Set[int]] = {} + # Deferred KV release: room -> current generation and ranks that + # ACKed. The generation rejects an old request's late ACK after room + # reuse. + self._deferred_abort_ack_tracker: Dict[int, DeferredAbortAckState] = {} + self._deferred_abort_generation = 0 # Heartbeat interval should be at least 2 seconds self.heartbeat_interval = max( envs.SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL.get(), 2.0 @@ -419,6 +533,11 @@ def update_status(self, bootstrap_room: int, status: KVPoll): # Anything else would resurrect a concluded room and pollute a later # request that reuses the same bootstrap_room. if status in (KVPoll.Bootstrapping, KVPoll.WaitingForInput): + # A new room lifecycle drops deferred-ACK state left by the + # previous request that used this bootstrap_room. + self._deferred_ack_targets.pop(bootstrap_room, None) + self._deferred_ack_poisoned_rooms.discard(bootstrap_room) + self._staging_outstanding.pop(bootstrap_room, None) self.request_status[bootstrap_room] = status return if status == KVPoll.Failed: @@ -638,24 +757,47 @@ def apply_prefill_status( prefill_rank, ) - def register_deferred_abort_room(self, bootstrap_room: int) -> None: - """Arm drain-ack accounting for a held room; a fresh set wipes stale acks - from a prior request that reused this bootstrap_room.""" - self._deferred_abort_ack_tracker[bootstrap_room] = set() + def register_deferred_abort_room(self, bootstrap_room: int) -> int: + """Arm a new generation for a held room and return its wire id.""" + self._deferred_abort_generation += 1 + generation = self._deferred_abort_generation + self._deferred_abort_ack_tracker[bootstrap_room] = DeferredAbortAckState( + generation=generation, + prefill_ranks=set(), + ) + return generation - def note_abort_ack(self, bootstrap_room: int, prefill_rank: int) -> None: - """Record a prefill rank's drain ack (decode receiver thread). Only counts - while the room is held; grabs the set by reference to avoid racing clear.""" - acks = self._deferred_abort_ack_tracker.get(bootstrap_room) - if acks is not None: - acks.add(prefill_rank) + def note_abort_ack( + self, bootstrap_room: int, prefill_rank: int, generation: int + ) -> None: + """Record an ACK only for the generation currently holding the room.""" + state = self._deferred_abort_ack_tracker.get(bootstrap_room) + if state is not None and state.generation == generation: + state.prefill_ranks.add(prefill_rank) + + def handle_abort_ack_message(self, msg: List[bytes]) -> bool: + """Consume an ABORT_ACK and record it for an armed decode room.""" + if not msg or msg[0] != ABORT_ACK_TAG: + return False + if not self.enable_deferred_decode_kv_release: + return True + ack = AbortAck.from_zmq(msg) + if ack is None: + return True + if ack.generation is None: + logger.warning_once( + "Generation-less ABORT_ACK received; ignoring it. " + "Mixed SGLang versions may wait for the KV release timeout." + ) + return True + self.note_abort_ack(ack.room, ack.prefill_rank, ack.generation) + return True def is_abort_release_safe(self, bootstrap_room: int, required_acks: int) -> bool: """True once every prefill rank that could still write these pages has acked.""" - return ( - len(self._deferred_abort_ack_tracker.get(bootstrap_room, ())) - >= required_acks - ) + state = self._deferred_abort_ack_tracker.get(bootstrap_room) + acks = state.prefill_ranks if state is not None else () + return len(acks) >= required_acks def clear_deferred_abort_state(self, bootstrap_room: int) -> None: self._deferred_abort_ack_tracker.pop(bootstrap_room, None) @@ -668,73 +810,55 @@ def _prefill_unique_rank(self) -> int: + self.attn_cp_rank ) - def _send_abort_ack(self, decode_ip: str, decode_port: int, room: int) -> None: + def _send_abort_ack(self, room: int, target: AckTarget) -> None: """Best-effort ack that this rank's transfer for an aborted room drained.""" try: - na = NetworkAddress(decode_ip, decode_port) + na = NetworkAddress(target.ip, target.port) self._send_multipart_locked( na.to_tcp(), - [ - b"ABORT_ACK", - str(room).encode("ascii"), - str(self._prefill_unique_rank()).encode("ascii"), - ], + AbortAck(room, self._prefill_unique_rank(), target.generation).to_zmq(), is_ipv6=na.is_ipv6, ) except Exception as e: logger.debug(f"Failed to send drained ABORT_ACK for room {room}: {e}") - def _abort_ack_fanout_targets(self, room: int) -> List[Tuple[str, int]]: - """Every decode peer of the room, dummy pairings included: each decode - rank counts drain acks from every prefill rank it notified, and a dummy - pairing still expects its ack (nothing was written, so it is trivially - drained). Snapshot: the control thread can register a late peer while - we walk the dict.""" - infos = self.transfer_infos.get(room) - if not infos: - return [] - targets: List[Tuple[str, int]] = [] - for info in list(infos.values()): - target = (info.endpoint, info.dst_port) - if target not in targets: - targets.append(target) - return targets - def _maybe_ack_drained_abort(self, room: int) -> None: """Send the deferred ack once an aborted room's chunks have drained (outstanding == 0). pop() makes it fire at most once.""" if self._staging_outstanding.get(room, 0) > 0: return - target = self._deferred_ack_targets.pop(room, None) - if target is None: - return - # With prefill TP < decode TP several decode ranks share one room, but - # the registry keeps only the last ABORT sender -- the earlier ranks - # would hold their pages for the full release timeout. Fan the ack out - # to every known peer: the live transfer_infos (catches peers that - # registered after the ABORT), the snapshot taken at registration - # (survives a teardown that popped transfer_infos mid-flight), and the - # registered sender as the guaranteed floor. - targets = self._abort_ack_fanout_targets(room) - for peer in self._deferred_ack_fanout_snapshots.pop(room, ()): - if peer not in targets: - targets.append(peer) - if target not in targets: - targets.append(target) - for decode_ip, decode_port in targets: - self._send_abort_ack(decode_ip, decode_port, room) - - def register_deferred_ack_target( - self, room: int, decode_ip: str, decode_port: int + self._send_abort_acks(room, self._deferred_ack_targets.pop(room, None)) + + def _send_abort_acks( + self, room: int, targets: Optional[Dict[Tuple[str, int], AckTarget]] ) -> None: - """Hold this room's ack until its transfer drains. Callers must mark the - room Failed FIRST -- registering while it still accepts chunks lets the - worker ack, then a new chunk writes pages the decode already released.""" - # Snapshot before target: the worker pops the target first, so writing - # the target first would let a drain racing this registration consume - # the target and strand the snapshot written after it. - self._deferred_ack_fanout_snapshots[room] = self._abort_ack_fanout_targets(room) - self._deferred_ack_targets[room] = (decode_ip, decode_port) + # With prefill TP < decode TP several decode ranks share one room and + # each sends its own ABORT; fan the ACK out to every one of them. + for target in (targets or {}).values(): + self._send_abort_ack(room, target) + + def register_deferred_ack_target(self, room: int, target: AckTarget) -> None: + """Hold this generation's ACK until the room's transfer drains. + + Targets are kept per decode endpoint, so each decode rank that aborted + the room gets an ACK for its own latest generation. Active-room callers + must mark the room Failed before registering so no new chunk can be + accepted after the ACK is sent. + """ + if room in self._deferred_ack_poisoned_rooms: + return + # Replace instead of mutating in place: a concurrent drain pops the dict + # and iterates it. If the pop wins, the caller's immediate + # _maybe_ack_drained_abort retry re-sends the merged set, and a + # duplicate ACK is harmless because decode records ACKs in a set. + targets = dict(self._deferred_ack_targets.get(room, ())) + targets[(target.ip, target.port)] = target + self._deferred_ack_targets[room] = targets + + def poison_deferred_ack_room(self, room: int) -> None: + """Drop an ACK target when the backend cannot prove the room drained.""" + self._deferred_ack_targets.pop(room, None) + self._deferred_ack_poisoned_rooms.add(room) def get_kv_replica_factor(self) -> int: if self._kv_replica_factor is None: @@ -1712,20 +1836,11 @@ def _check_bootstrap_timeout(self) -> Optional[KVPoll]: def clear(self) -> None: self.kv_mgr.request_status.pop(self.bootstrap_room, None) - if hasattr(self.kv_mgr, "req_to_decode_prefix_len"): - self.kv_mgr.req_to_decode_prefix_len.pop(self.bootstrap_room, None) - if hasattr(self.kv_mgr, "transfer_infos"): - self.kv_mgr.transfer_infos.pop(self.bootstrap_room, None) - if hasattr(self.kv_mgr, "_deferred_ack_targets"): - if hasattr(self.kv_mgr, "_staging_outstanding"): - # Preserve the target until in-flight writes drain, even when - # the scheduler has already observed Failed and cleared the room. - self.kv_mgr._maybe_ack_drained_abort(self.bootstrap_room) - else: - self.kv_mgr._deferred_ack_targets.pop(self.bootstrap_room, None) - self.kv_mgr._deferred_ack_fanout_snapshots.pop( - self.bootstrap_room, None - ) + self.kv_mgr.req_to_decode_prefix_len.pop(self.bootstrap_room, None) + self.kv_mgr.transfer_infos.pop(self.bootstrap_room, None) + # A sender can be cleared while a counted transfer is still writing: + # keep the target for the worker, or ACK now if nothing is in flight. + self.kv_mgr._maybe_ack_drained_abort(self.bootstrap_room) def abort(self): self.kv_mgr.record_failure( @@ -1756,6 +1871,7 @@ def __init__( self.require_staging: bool = False self.init_time: Optional[float] = None self.abort_notified: bool = False + self._abort_generation: Optional[int] = None self._connection_pool_entries: Dict[str, List[Dict]] = {} self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping) @@ -2054,30 +2170,32 @@ def ensure_abort_notified(self, *, force_arm: bool = False) -> None: self.abort_notified = True def _send_abort_notification(self, *, force_arm: bool = False): - # Once metadata is published (init_time set) prefill may already be - # writing; arm the drain-ack tracker BEFORE the ABORT goes out, so an - # ack racing back -- or fanned out by a peer rank's earlier abort of - # the same room -- is counted instead of dropped. Prealloc-queue - # receivers (init_time None) never enter the deferred-release flow - # that would clean the tracker up, so they stay unarmed -- except on a - # partial publish, where init_time is still None but earlier ranks - # already hold destinations; those callers defer and pass force_arm. + # Both abort() and the waiting timeout reach here. Once metadata is + # published (init_time set) prefill may be writing, so arm the ACK + # tracker before the ABORT goes out and an immediate ACK can land. + # Prealloc-queue receivers (init_time None) never enter the + # deferred-release flow that would clean the tracker up, so they stay + # unarmed -- except on a partial publish, where init_time is still None + # but earlier ranks already hold destinations; those callers defer and + # pass force_arm. if self.kv_mgr.enable_deferred_decode_kv_release and ( force_arm or self.init_time is not None ): - self.kv_mgr.register_deferred_abort_room(self.bootstrap_room) + self._abort_generation = self.kv_mgr.register_deferred_abort_room( + self.bootstrap_room + ) for bootstrap_info in self.bootstrap_infos: # Best-effort notification to prefill side that this request was aborted. try: sock, lock = self._connect_to_bootstrap_server(bootstrap_info) with lock: sock.send_multipart( - [ - b"ABORT", - str(self.bootstrap_room).encode("ascii"), - self.kv_mgr.local_ip.encode("ascii"), - str(self.kv_mgr.rank_port).encode("ascii"), - ] + AbortNotification( + room=self.bootstrap_room, + decode_ip=self.kv_mgr.local_ip, + decode_port=self.kv_mgr.rank_port, + generation=self._abort_generation, + ).to_zmq() ) logger.debug( f"Sent abort notification for room {self.bootstrap_room} " diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index e60b437a2fb3..9f1ade96d844 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -17,6 +17,9 @@ from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll, StateType from sglang.srt.disaggregation.common.conn import ( + ABORT_ACK_TAG, + ABORT_TAG, + AbortNotification, CommonKVBootstrapServer, CommonKVManager, CommonKVReceiver, @@ -249,9 +252,6 @@ def __init__( self.failed_sessions = set() self.session_lock = threading.Lock() self.start_prefill_thread() - # Per-room count of chunks not yet transferred; teardown waits for - # zero so a deferred chunk is not dropped by an early conclude. - self._staging_outstanding = defaultdict(int) # Determine the number of threads to use for kv sender cpu_count = os.cpu_count() transfer_thread_pool_size = ( @@ -2134,8 +2134,9 @@ def transfer_worker( ) while True: + kv_chunk: Optional[TransferKVChunk] = None try: - kv_chunk: TransferKVChunk = queue.get() + kv_chunk = queue.get() # teardown() pushes a None sentinel to unblock get() and stop # the worker: FastQueue.get() blocks indefinitely, so checking # _stopped alone can never wake a parked worker during a role switch. @@ -2470,10 +2471,83 @@ def transfer_worker( except Exception as e: # NOTE(shangming): Remove this when we make sure the transfer thread is bug-free + if kv_chunk is not None: + self.poison_deferred_ack_room(kv_chunk.room) raise RuntimeError( f"Transfer thread failed because of {e}. Prefill instance with bootstrap_port={self.bootstrap_port} is dead." ) + def _handle_abort_notification(self, msg: List[bytes]) -> bool: + if not msg or msg[0] != ABORT_TAG: + return False + + notification = AbortNotification.from_zmq(msg) + if notification is None: + return True + room = notification.room + status = self.request_status.get(room) + room_active = status is not None and status != KVPoll.Success + + if self.enable_deferred_decode_kv_release: + self._handle_deferred_abort_notification(notification, room_active) + else: + self._handle_legacy_abort_notification(notification, room_active) + return True + + def _handle_deferred_abort_notification( + self, notification: AbortNotification, room_active: bool + ) -> None: + room = notification.room + if room_active: + self.update_status(room, KVPoll.Failed) + + ack_target = notification.deferred_ack_target() + if ack_target is None: + return + self.register_deferred_ack_target(room, ack_target) + self._maybe_ack_drained_abort(room) + if room_active: + logger.debug( + "Received abort notification for room %s, marked as Failed; " + "ACK deferred until transfer drains", + room, + ) + + def _handle_legacy_abort_notification( + self, notification: AbortNotification, room_active: bool + ) -> None: + room = notification.room + if room_active: + self.update_status(room, KVPoll.Failed) + logger.debug( + "Received abort notification for room %s, marked as Failed", room + ) + else: + logger.debug( + "Received abort notification for room %s, ignoring " + "(already completed or unknown)", + room, + ) + + if notification.decode_ip is None or notification.decode_port is None: + return + try: + na = NetworkAddress(notification.decode_ip, notification.decode_port) + # Deferred ACKs carry a rank; the feature-off wire format does not. + self._send_multipart_locked( + na.to_tcp(), + [ABORT_ACK_TAG, str(room).encode("ascii")], + is_ipv6=na.is_ipv6, + ) + logger.debug( + "Sent ABORT_ACK for room %s to %s:%s", + room, + notification.decode_ip, + notification.decode_port, + ) + except Exception as exc: + logger.debug("Failed to send ABORT_ACK for room %s: %s", room, exc) + def start_prefill_thread(self): recv = self._make_worker_recv(self.server_socket) @@ -2493,76 +2567,7 @@ def bootstrap_thread(): if room == "STAGING_RSP": handle_staging_rsp(waiting_req_bytes, self.transfer_infos) continue - # Decode-side abort notification: mark room as failed and ACK - if room == "ABORT": - room_to_be_aborted = int(waiting_req_bytes[1].decode("ascii")) - decode_ip = waiting_req_bytes[2].decode("ascii") - decode_port = int(waiting_req_bytes[3].decode("ascii")) - room_active = ( - room_to_be_aborted in self.request_status - and self.check_status(room_to_be_aborted) != KVPoll.Success - ) - if self.enable_deferred_decode_kv_release: - # Mark Failed FIRST (stops add_transfer_request enqueuing - # new chunks), THEN register the ack target: registering - # first would let the worker drain+ack while the room is - # not yet Failed, so a newly enqueued chunk could still - # write to the freed pages. The worker (not this thread) - # acks once its in-flight write drains. - if ( - room_active - or self._staging_outstanding.get(room_to_be_aborted, 0) > 0 - ): - self.update_status(room_to_be_aborted, KVPoll.Failed) - self.register_deferred_ack_target( - room_to_be_aborted, decode_ip, decode_port - ) - # Try once: the room may already be quiescent and - # never revisited by the worker. - self._maybe_ack_drained_abort(room_to_be_aborted) - logger.debug( - f"Received abort notification for room {room_to_be_aborted}, " - f"marked as Failed; ACK deferred until transfer drains" - ) - else: - # Concluded/unknown AND quiescent (the branch above - # already took every case with writes outstanding): - # ack now so decode releases without the timeout. - self._send_abort_ack( - decode_ip, decode_port, room_to_be_aborted - ) - continue - # No need to abort the room if it has already succeeded - if room_active: - self.update_status(room_to_be_aborted, KVPoll.Failed) - logger.debug( - f"Received abort notification for room {room_to_be_aborted}, " - f"marked as Failed" - ) - else: - logger.debug( - f"Received abort notification for room {room_to_be_aborted}, " - f"ignoring (already completed or unknown)" - ) - # Send ACK back to decode endpoint - try: - na = NetworkAddress(decode_ip, decode_port) - self._send_multipart_locked( - na.to_tcp(), - [ - b"ABORT_ACK", - str(room_to_be_aborted).encode("ascii"), - ], - is_ipv6=na.is_ipv6, - ) - logger.debug( - f"Sent ABORT_ACK for room {room_to_be_aborted} to " - f"{decode_ip}:{decode_port}" - ) - except Exception as e: - logger.debug( - f"Failed to send ABORT_ACK for room {room_to_be_aborted}: {e}" - ) + if self._handle_abort_notification(waiting_req_bytes): continue mooncake_session_id = waiting_req_bytes[3].decode("ascii") if room == "None": @@ -2659,16 +2664,7 @@ def decode_thread(): self._handle_staging_req(msg) continue - # Prefill acknowledges abort notification - if msg[0] == b"ABORT_ACK": - ack_aborted_room = int(msg[1].decode("ascii")) - logger.debug(f"Received ABORT_ACK for room {ack_aborted_room}") - # Deferred release: the 3-frame ack carries the prefill rank - # and means its transfer drained; aggregate for is_abort_release_safe. - if self.enable_deferred_decode_kv_release and len(msg) >= 3: - self.note_abort_ack( - ack_aborted_room, int(msg[2].decode("ascii")) - ) + if self.handle_abort_ack_message(msg): continue parsed = self.parse_kv_status_message(msg) diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 475f38bcb6e6..370f38012e73 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -7,7 +7,8 @@ import threading import time import uuid -from typing import List, Optional +from queue import Full, Queue +from typing import List, Optional, Tuple import msgspec import numpy as np @@ -28,6 +29,9 @@ from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll from sglang.srt.disaggregation.common.conn import ( + ABORT_TAG, + AbortNotification, + AckTarget, CommonKVBootstrapServer, CommonKVManager, CommonKVReceiver, @@ -50,7 +54,6 @@ logger = logging.getLogger(__name__) MORI_GUARD = b"MoriMsgGuard" -_TAG_ABORT = b"ABORT" def _normalize_state_indices_per_component( @@ -291,10 +294,28 @@ class TransferTarget: peer_info: KVArgsRegisterInfo +class _MoriTransferSubmissionError(RuntimeError): + def __init__(self, message: str, statuses: List[TransferStatus]): + super().__init__(message) + self.statuses = statuses + + +class _SubmissionLocal(threading.local): + """Per-thread sink for statuses issued by a submission still in progress. + + `statuses` is None outside a tracked submission, so a leaf that records + into it is a no-op on threads that are not submitting. + """ + + statuses: Optional[List[TransferStatus]] = None + + class MoriKVManager(CommonKVManager): AUX_DATA_HEADER = b"AUX_DATA" # Implements teardown() below, so runtime PD role switching is supported. supports_role_switch = True + # The ABORT ack is held until the transfer worker or drainer quiesces the room. + supports_deferred_decode_kv_release = True # The bootstrap socket carries several message kinds, so the status message # is tagged. Mori has always shipped the failure reason with it. @@ -315,6 +336,7 @@ def __init__( self.aux_mem_descs: List[MemoryDesc] = [] self.state_mem_descs: List[List[MemoryDesc]] = [] self.transfer_lock = threading.Lock() + self._submission_local = _SubmissionLocal() self._zmq_ctx = zmq.Context() self._socket_local = threading.local() # Set by teardown() to make worker threads exit (PoC: P<->D role switch). @@ -329,6 +351,11 @@ def __init__( ] self._wait_poll_ms = envs.SGLANG_MORI_WAIT_POLL_MS.get() self._transfer_timeout_ms = envs.SGLANG_MORI_TRANSFER_TIMEOUT_MS.get() + self._abort_ack_lock = threading.Lock() + if self.enable_deferred_decode_kv_release: + self._drain_queue: Queue[ + Optional[Tuple[TransferKVChunk, List[TransferStatus]]] + ] = Queue(maxsize=self._num_shards) for shard, queue in enumerate(self._transfer_queues): # Track the thread so teardown() can join it: otherwise every # P->D->P flip that re-enters PREFILL leaks _num_shards threads @@ -344,6 +371,17 @@ def __init__( ) t.start() self._worker_threads.append(t) + if self.enable_deferred_decode_kv_release: + t = threading.Thread( + target=self._drain_worker, + daemon=True, + name=( + f"mori-drain-dp{self.system_dp_rank}-" + f"tp{self.attn_tp_rank}-s{shard}" + ), + ) + t.start() + self._worker_threads.append(t) self._start_bootstrap_thread() elif self.disaggregation_mode == DisaggregationMode.DECODE: self._start_decode_thread() @@ -435,8 +473,16 @@ def _transfer_worker(self, queue: FastQueue) -> None: # alone can never wake a parked worker during a role switch. if kv_chunk is None: break + self._mark_transfer_started(kv_chunk) try: - self._process_transfer_chunk(kv_chunk) + is_quiescent = self._process_transfer_chunk(kv_chunk) + except _MoriTransferSubmissionError as exc: + logger.exception( + "Mori transfer submission failed for room %s", + kv_chunk.room, + ) + self._handle_submission_failure(kv_chunk, exc) + continue except Exception as exc: failure_reason = f"transfer worker raised: {exc!r}" try: @@ -458,45 +504,174 @@ def _transfer_worker(self, queue: FastQueue) -> None: ) except Exception: pass + self._mark_transfer_quiescent(kv_chunk) + continue + + if is_quiescent: + self._mark_transfer_quiescent(kv_chunk) - def _process_transfer_chunk(self, kv_chunk: TransferKVChunk) -> None: + def _handle_submission_failure( + self, kv_chunk: TransferKVChunk, exc: _MoriTransferSubmissionError + ) -> None: + failure_reason = str(exc) + if exc.statuses: + self._enqueue_transfer_drain(kv_chunk, exc.statuses, failure_reason) + return + self._mark_transfer_quiescent(kv_chunk) + self.conclude_failure( + bootstrap_room=kv_chunk.room, failure_reason=failure_reason + ) + + def _mark_transfer_started(self, kv_chunk: TransferKVChunk) -> None: + if not self.enable_deferred_decode_kv_release: + return + with self._abort_ack_lock: + self._staging_outstanding[kv_chunk.room] += 1 + + def _mark_transfer_quiescent(self, kv_chunk: TransferKVChunk) -> None: + if not self.enable_deferred_decode_kv_release: + return + + room = kv_chunk.room + targets = None + with self._abort_ack_lock: + self._staging_outstanding[room] -= 1 + if self._staging_outstanding[room] <= 0: + self._staging_outstanding.pop(room) + targets = self._deferred_ack_targets.pop(room, None) + self._send_abort_acks(room, targets) + + def _process_transfer_chunk(self, kv_chunk: TransferKVChunk) -> bool: room = kv_chunk.room if self._should_skip_transfer(room): - return + return True if kv_chunk.wait_event is not None: kv_chunk.wait_event.synchronize() if self._should_skip_transfer(room): - return + return True - statuses = self._submit_kv_transfer( - room, - kv_chunk.prefill_kv_indices, - kv_chunk.index_slice, - kv_chunk.is_last_chunk, - aux_index=kv_chunk.prefill_aux_index, - state_indices=kv_chunk.state_indices, - ) + statuses = self._submit_transfer_chunk(kv_chunk) if self._should_skip_transfer(room): - return + if not self.enable_deferred_decode_kv_release: + return True + # Aborted after submit: the writes must still quiesce before the ACK, + # but that wait belongs to the drainer, not this shard worker. + self._enqueue_transfer_drain(kv_chunk, statuses, None) + return False - failure_reason = self._wait_transfer_completion(statuses) + failure_reason, is_quiescent = self._wait_for_chunk_completion( + kv_chunk, statuses + ) + if not is_quiescent: + # The drainer now owns the decrement. + return False if self._should_skip_transfer(room): - return + return True if failure_reason is not None: self.conclude_failure(bootstrap_room=room, failure_reason=failure_reason) - return + return True if kv_chunk.is_last_chunk: # conclude_transfer downgrades to Failed when a failure was recorded # while this chunk was in flight, and applies the same status locally # and on the wire. self.conclude_transfer(bootstrap_room=room, status=KVPoll.Success) + return True + + def _wait_for_chunk_completion( + self, kv_chunk: TransferKVChunk, statuses: List[TransferStatus] + ) -> Tuple[Optional[str], bool]: + try: + failure_reason = self._wait_transfer_completion(statuses) + is_quiescent = not self.enable_deferred_decode_kv_release or not any( + status.InProgress() for status in statuses + ) + except Exception as exc: + if not self.enable_deferred_decode_kv_release: + raise + failure_reason = f"Transfer completion failed: {exc!r}" + is_quiescent = False + + if not is_quiescent: + self._enqueue_transfer_drain(kv_chunk, statuses, failure_reason) + return failure_reason, is_quiescent + + def _submit_transfer_chunk(self, kv_chunk: TransferKVChunk) -> List[TransferStatus]: + if not self.enable_deferred_decode_kv_release: + return self._submit_kv_transfer( + kv_chunk.room, + kv_chunk.prefill_kv_indices, + kv_chunk.index_slice, + kv_chunk.is_last_chunk, + aux_index=kv_chunk.prefill_aux_index, + state_indices=kv_chunk.state_indices, + ) + + self._submission_local.statuses = [] + try: + return self._submit_kv_transfer( + kv_chunk.room, + kv_chunk.prefill_kv_indices, + kv_chunk.index_slice, + kv_chunk.is_last_chunk, + aux_index=kv_chunk.prefill_aux_index, + state_indices=kv_chunk.state_indices, + ) + except Exception as exc: + raise _MoriTransferSubmissionError( + f"Transfer submission failed: {exc}", + list(self._submission_local.statuses), + ) from exc + finally: + self._submission_local.statuses = None + + def _enqueue_transfer_drain( + self, + kv_chunk: TransferKVChunk, + statuses: List[TransferStatus], + failure_reason: Optional[str], + ) -> None: + # Conclude before the handoff so decode hears now and later chunks skip; + # only the ACK waits for the drain. + if failure_reason is not None: + self.conclude_failure( + bootstrap_room=kv_chunk.room, failure_reason=failure_reason + ) + self._drain_queue.put((kv_chunk, statuses)) + + def _drain_worker(self) -> None: + while True: + item = self._drain_queue.get() + # teardown() pushes a None sentinel, as for the transfer workers. + if item is None: + break + kv_chunk, statuses = item + while True: + try: + self._drain_transfer_statuses(kv_chunk, statuses) + break + except Exception: + logger.exception( + "Mori transfer drainer failed for room %s; retrying", + kv_chunk.room, + ) + time.sleep(max(self._wait_poll_ms, 1) / 1000) + + def _drain_transfer_statuses( + self, kv_chunk: TransferKVChunk, statuses: List[TransferStatus] + ) -> None: + for status in statuses: + if status.InProgress(): + status.Wait() + self._mark_transfer_quiescent(kv_chunk) def _should_skip_transfer(self, room: int) -> bool: - if room not in self.request_status or self.check_status(room) == KVPoll.Failed: + # Single read: clear() may pop the room concurrently. + status = self.request_status.get(room) + if status is None or status == KVPoll.Failed: logger.debug( "Skipping chunk for room %s because it has already failed or been aborted", room, @@ -665,37 +840,44 @@ def _validate_message(self, msg: List[bytes]) -> Optional[List[bytes]]: return payload def _handle_abort_message(self, msg: List[bytes]) -> None: - """Handle best-effort ABORT notifications from the decode side.""" - if len(msg) < 2: - logger.warning("Malformed ABORT message: too few frames (%d)", len(msg)) - return - - try: - bootstrap_room = int(msg[1].decode("ascii")) - except (ValueError, UnicodeDecodeError): - logger.warning("Malformed ABORT message: invalid room field %r", msg[1]) + """Handle ABORT and defer its ACK until this rank's writes drain.""" + notification = AbortNotification.from_zmq(msg) + if notification is None: return + bootstrap_room = notification.room with self.transfer_lock: current = self.request_status.get(bootstrap_room) - if current is None: - logger.debug( - "ABORT for room %s is not tracked; ignoring", - bootstrap_room, - ) - return - if current == KVPoll.Success: - logger.debug( - "ABORT for room %s already succeeded; ignoring", - bootstrap_room, - ) - return - if current == KVPoll.Failed: - return + room_active = current is not None and current != KVPoll.Success + if room_active and current != KVPoll.Failed: + self.update_status(bootstrap_room, KVPoll.Failed) + + if room_active: + logger.debug("Room %s marked Failed via ABORT from decode", bootstrap_room) + else: + logger.debug( + "ABORT for room %s is already complete or not tracked", + bootstrap_room, + ) - self.update_status(bootstrap_room, KVPoll.Failed) + self._arm_abort_ack(bootstrap_room, notification) + + def _arm_abort_ack( + self, + bootstrap_room: int, + notification: AbortNotification, + ) -> None: + if not self.enable_deferred_decode_kv_release: + return + ack_target: Optional[AckTarget] = notification.deferred_ack_target() + if ack_target is None: + return - logger.debug("Room %s marked Failed via ABORT from decode", bootstrap_room) + with self._abort_ack_lock: + if self._staging_outstanding.get(bootstrap_room, 0) > 0: + self.register_deferred_ack_target(bootstrap_room, ack_target) + return + self._send_abort_ack(bootstrap_room, ack_target) def _start_bootstrap_thread(self) -> None: recv = self._make_worker_recv(self.server_socket) @@ -708,7 +890,7 @@ def bootstrap_worker(): continue tag = msg[0] - if tag == _TAG_ABORT: + if tag == ABORT_TAG: self._handle_abort_message(msg) continue @@ -742,6 +924,8 @@ def decode_worker(): if msg and msg[0] == MoriKVManager.AUX_DATA_HEADER: self._handle_aux_data(msg) continue + if self.handle_abort_ack_message(msg): + continue parsed = self.parse_kv_status_message(msg) if parsed is None: @@ -781,6 +965,17 @@ def teardown(self) -> None: queue.put(None) except Exception: logger.exception("Failed to signal mori transfer worker on teardown") + # Drain workers park in Queue.get() the same way. A full queue means a + # drainer is stuck in status.Wait(); its join below times out. + if ( + self.disaggregation_mode == DisaggregationMode.PREFILL + and self.enable_deferred_decode_kv_release + ): + for _ in range(self._num_shards): + try: + self._drain_queue.put_nowait(None) + except Full: + break # Join workers before touching their sockets: ZMQ sockets aren't # thread-safe, so don't close server_socket while a worker may poll it. for t in self._worker_threads: @@ -922,8 +1117,14 @@ def _submit_batch_transfer_plan( [plan.sizes], [transfer_uid], ) + self._record_submitted_statuses(statuses) return statuses + def _record_submitted_statuses(self, statuses: List[TransferStatus]) -> None: + active_statuses = self._submission_local.statuses + if active_statuses is not None: + active_statuses.extend(statuses) + def _build_contiguous_transfer_plan( self, grouped_plan: GroupedIndexPlan, item_len: int ) -> BatchTransferPlan: @@ -1157,11 +1358,13 @@ def send_aux_rdma( remote_offsets.append([dst_aux_index * item_len]) sizes.append([item_len]) uids.append(self.engine.allocate_transfer_uid()) - return list( + statuses = list( self.engine.batch_write( src_descs, local_offsets, dst_descs, remote_offsets, sizes, uids ) ) + self._record_submitted_statuses(statuses) + return statuses def send_aux_tcp( self, @@ -1384,6 +1587,7 @@ def _send_mamba_state( [[size]], [transfer_uid], ) + self._record_submitted_statuses(batch_statuses) statuses.extend(batch_statuses) return statuses @@ -1540,46 +1744,37 @@ def _submit_kv_transfer( targets.append(TransferTarget(info=info, peer_info=peer_info)) result_statuses: List[TransferStatus] = [] - try: - for target in targets: - info = target.info - peer_info = target.peer_info - - if not info.is_dummy: - dst_indices_chunk = info.dst_kv_indices[index_slice] - result_statuses.extend( - self.send_kvcache(peer_info, kv_indices, dst_indices_chunk) - ) + for target in targets: + info = target.info + peer_info = target.peer_info + + if not info.is_dummy: + dst_indices_chunk = info.dst_kv_indices[index_slice] + result_statuses.extend( + self.send_kvcache(peer_info, kv_indices, dst_indices_chunk) + ) - if ( - is_last_chunk - and state_indices is not None - and not info.is_dummy - and self.state_mem_descs - ): - result_statuses.extend( - self.send_state( - peer_info, state_indices, info.dst_state_indices - ) - ) + if ( + is_last_chunk + and state_indices is not None + and not info.is_dummy + and self.state_mem_descs + ): + result_statuses.extend( + self.send_state(peer_info, state_indices, info.dst_state_indices) + ) - if ( - is_last_chunk - and aux_index is not None - and info.dst_aux_index >= 0 - and self.pp_group.is_last_rank - ): - result_statuses.extend( - self.send_aux( - peer_info, aux_index, info.dst_aux_index, bootstrap_room - ) + if ( + is_last_chunk + and aux_index is not None + and info.dst_aux_index >= 0 + and self.pp_group.is_last_rank + ): + result_statuses.extend( + self.send_aux( + peer_info, aux_index, info.dst_aux_index, bootstrap_room ) - except Exception as e: - logger.exception( - "Mori KV transfer submission failed for bootstrap_room=%s", - bootstrap_room, - ) - raise RuntimeError(f"Transfer submission failed: {e}") from e + ) return result_statuses @@ -1670,6 +1865,13 @@ def poll(self) -> KVPoll: self.conclude_state = status return status + def clear(self) -> None: + if self.kv_mgr.enable_deferred_decode_kv_release: + with self.kv_mgr._abort_ack_lock: + super().clear() + else: + super().clear() + def failure_exception(self): if self.conclude_state is None: self.conclude_state = KVPoll.Failed diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 20a6bc7b12c3..7e42c2039fee 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -20,6 +20,8 @@ from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll, StateType from sglang.srt.disaggregation.common.conn import ( + ABORT_TAG, + AbortNotification, CommonKVBootstrapServer, CommonKVManager, CommonKVReceiver, @@ -536,9 +538,6 @@ def create_backend(): FastQueue() for _ in range(transfer_queue_size) ] self.exceptions: Dict[int, Exception] = {} - # Per-room count of chunks not yet transferred; teardown waits for - # zero so a deferred chunk is not dropped by an early conclude. - self._staging_outstanding = defaultdict(int) # Mirror mooncake: one staging buffer per worker queue, all # built before workers spawn so each worker owns a private # buffer (no cross-worker contention on the staging ring). @@ -647,12 +646,7 @@ def decode_listener_thread(): if self.enable_staging: self._handle_staging_req(msg) continue - if msg[0] == b"ABORT_ACK": - # Drain ack for an aborted room; aggregate per prefill rank. - if len(msg) >= 3: - self.note_abort_ack( - int(msg[1].decode("ascii")), int(msg[2].decode("ascii")) - ) + if self.handle_abort_ack_message(msg): continue parsed = self.parse_kv_status_message(msg) if parsed is not None: @@ -1238,6 +1232,9 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) room, ) self._staging_outstanding.pop(room, None) + if self.enable_deferred_decode_kv_release: + # clear() keeps the target while a chunk is counted. + self._maybe_ack_drained_abort(room) continue # Counted at dequeue, before the status check, so @@ -1263,6 +1260,8 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) room, ) self._staging_outstanding.pop(room, None) + if self.enable_deferred_decode_kv_release: + self._maybe_ack_drained_abort(room) continue # Lazily build a per-worker staging strategy bound to this @@ -1572,6 +1571,10 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) # than telling it those pages are free. self.record_failure(room, str(e)) self.update_status(room, KVPoll.Failed) + # This chunk stays counted in _staging_outstanding, so no + # drain ACK can follow; discard the target and fall back to + # the timeout. + self.poison_deferred_ack_room(room) def register_buffer_to_engine(self): self.kv_descs = [] @@ -3004,16 +3007,13 @@ def check_transfer_done(self, room: int): return self.transfer_statuses[room].is_done() def _handle_abort_notification(self, msg: List[bytes]) -> bool: - if not msg or msg[0] != b"ABORT": + if not msg or msg[0] != ABORT_TAG: return False - try: - room_to_be_aborted = int(msg[1].decode("ascii")) - decode_ip = msg[2].decode("ascii") if len(msg) > 2 else None - decode_port = int(msg[3].decode("ascii")) if len(msg) > 3 else None - except Exception as e: - logger.debug(f"Ignoring malformed abort notification: {e}") + notification = AbortNotification.from_zmq(msg) + if notification is None: return True + room_to_be_aborted = notification.room room_active = ( room_to_be_aborted in self.request_status @@ -3035,23 +3035,24 @@ def _handle_abort_notification(self, msg: List[bytes]) -> bool: f"ignoring (already completed or unknown)" ) - # Deferred KV release: register only after the status flip above (see - # register_deferred_ack_target), then try once -- the room may already be - # quiescent and never revisited by the worker. A concluded/unknown room is - # acked only when nothing is still counted for it: the ERR path abandons - # sibling handles that may still be writing and clear() then drops the - # room, so "unknown" alone does not imply quiescent. - if self.enable_deferred_decode_kv_release and decode_port is not None: - if room_active: - self.register_deferred_ack_target( - room_to_be_aborted, decode_ip, decode_port - ) - self._maybe_ack_drained_abort(room_to_be_aborted) - elif self._staging_outstanding.get(room_to_be_aborted, 0) == 0: - self._send_abort_ack(decode_ip, decode_port, room_to_be_aborted) - + self._handle_deferred_abort_ack(notification) return True + def _handle_deferred_abort_ack(self, notification: AbortNotification) -> None: + room_to_be_aborted = notification.room + if not self.enable_deferred_decode_kv_release: + return + ack_target = notification.deferred_ack_target() + if ack_target is None: + return + + # The active-room status flip happens before registration. Success or a + # missing status does not imply quiescence: clear() can remove the room + # while a counted handle is still writing. The immediate retry closes + # both races where the worker drains before or during registration. + self.register_deferred_ack_target(room_to_be_aborted, ack_target) + self._maybe_ack_drained_abort(room_to_be_aborted) + def _start_bootstrap_thread(self): def bootstrap_thread(): """This thread recvs transfer info from the decode engine""" diff --git a/test/registered/unit/disaggregation/test_dcp_pack.py b/test/registered/unit/disaggregation/test_dcp_pack.py index dc5190d37224..3919c7e52a94 100644 --- a/test/registered/unit/disaggregation/test_dcp_pack.py +++ b/test/registered/unit/disaggregation/test_dcp_pack.py @@ -207,6 +207,9 @@ def test_cached_prefix_fits_pack_capacity_and_preserves_pages_and_state(self): mgr.transfer_queues = [None] mgr._register_staging_memory = Mock() mgr.request_status = {} + mgr._deferred_ack_targets = {} + mgr._deferred_ack_poisoned_rooms = set() + mgr._staging_outstanding = {} mgr.is_dummy_cp_rank = False mgr.enable_all_cp_ranks_for_transfer = False mgr.decode_kv_args_table = { diff --git a/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py b/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py index de45d80a6ad0..21759941c27a 100644 --- a/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py +++ b/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py @@ -11,11 +11,21 @@ import threading import unittest from types import SimpleNamespace +from typing import NamedTuple from unittest.mock import patch from sglang.srt.disaggregation import decode as decode_mod from sglang.srt.disaggregation.base.conn import BaseKVManager, KVPoll -from sglang.srt.disaggregation.common.conn import CommonKVManager, CommonKVReceiver +from sglang.srt.disaggregation.common.conn import ( + ABORT_ACK_TAG, + ABORT_TAG, + AbortAck, + AbortNotification, + AckTarget, + CommonKVManager, + CommonKVReceiver, + CommonKVSender, +) from sglang.srt.disaggregation.decode import DecodeTransferQueue from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager from sglang.srt.disaggregation.nixl.conn import NixlKVManager @@ -25,49 +35,492 @@ register_cpu_ci(est_time=10, suite="base-a-test-cpu") +ABORT_GENERATION = 7 + + +class AbortScenario(NamedTuple): + name: str + status: KVPoll | None + outstanding: int + expected_status: KVPoll | None + + +ABORT_SCENARIOS = ( + AbortScenario( + name="active transfer", + status=KVPoll.Transferring, + outstanding=1, + expected_status=KVPoll.Failed, + ), + # Window 2: no worker will revisit a room that is already quiescent. + AbortScenario( + name="active quiescent room", + status=KVPoll.WaitingForInput, + outstanding=0, + expected_status=KVPoll.Failed, + ), + AbortScenario( + name="completed transfer", + status=KVPoll.Success, + outstanding=1, + expected_status=KVPoll.Success, + ), + AbortScenario( + name="completed quiescent room", + status=KVPoll.Success, + outstanding=0, + expected_status=KVPoll.Success, + ), + # clear() can remove request_status before a counted write drains. + AbortScenario( + name="untracked transfer", + status=None, + outstanding=1, + expected_status=None, + ), + AbortScenario( + name="untracked quiescent room", + status=None, + outstanding=0, + expected_status=None, + ), +) + def _make_manager(): """A bare CommonKVManager carrying only the deferred-ack state the helpers touch (avoids the heavy real __init__).""" mgr = CommonKVManager.__new__(CommonKVManager) mgr._deferred_abort_ack_tracker = {} + mgr._deferred_abort_generation = 0 + mgr.enable_deferred_decode_kv_release = True + return mgr + + +def _make_prefill_manager(): + mgr = CommonKVManager.__new__(CommonKVManager) + mgr.enable_deferred_decode_kv_release = True + mgr.request_status = {} + mgr.req_to_decode_prefix_len = {} + mgr.transfer_infos = {} + mgr._deferred_ack_targets = {} + mgr._deferred_ack_poisoned_rooms = set() + mgr._staging_outstanding = {} + mgr._sent = [] + mgr._send_abort_ack = lambda *args: mgr._sent.append(args) return mgr +class _TestReceiver(CommonKVReceiver): + def poll(self): + raise NotImplementedError + + def failure_exception(self): + raise NotImplementedError + + +class _TestSender(CommonKVSender): + def poll(self): + raise NotImplementedError + + def failure_exception(self): + raise NotImplementedError + + +class DeferredAbortNotificationScenarios: + room: int + decode_ip: str + decode_port: int + + def _make_abort_manager(self, status: KVPoll | None): + raise NotImplementedError + + def _dispatch_abort(self, manager) -> None: + claimed = manager._handle_abort_notification(self._abort_message()) + self.assertTrue(claimed) + + def _start_test_transfer(self, manager) -> None: + manager._staging_outstanding[self.room] = 1 + + def _drain_test_transfer(self, manager) -> None: + manager._staging_outstanding[self.room] -= 1 + manager._maybe_ack_drained_abort(self.room) + + def _abort_message(self) -> list[bytes]: + return AbortNotification( + self.room, + self.decode_ip, + self.decode_port, + ABORT_GENERATION, + ).to_zmq() + + def test_deferred_ack_follows_room_status_and_outstanding_transfers(self): + target = AckTarget( + self.decode_ip, + self.decode_port, + ABORT_GENERATION, + ) + for case in ABORT_SCENARIOS: + with self.subTest(name=case.name): + manager = self._make_abort_manager(case.status) + if case.outstanding: + self._start_test_transfer(manager) + + self._dispatch_abort(manager) + + self.assertEqual( + manager.request_status.get(self.room), case.expected_status + ) + if case.outstanding: + self.assertEqual( + manager._deferred_ack_targets[self.room], + {(target.ip, target.port): target}, + ) + self.assertEqual(manager._sent, []) + self._drain_test_transfer(manager) + self.assertNotIn(self.room, manager._deferred_ack_targets) + self.assertEqual(manager._sent, [(self.room, target)]) + + +class TaggedAbortNotificationScenarios: + room: int + + def test_non_abort_message_is_not_claimed(self): + manager = self._make_abort_manager(KVPoll.WaitingForInput) + + self.assertFalse( + manager._handle_abort_notification( + [b"STAGING_REQ", str(self.room).encode()] + ) + ) + + +class WorkerFailureAbortScenarios: + room: int + decode_ip: str + decode_port: int + + def _provoke_worker_failure(self, manager) -> None: + raise NotImplementedError + + def test_worker_exception_poison_is_reset_for_reused_room(self): + target = AckTarget( + self.decode_ip, + self.decode_port, + ABORT_GENERATION, + ) + manager = self._make_abort_manager(KVPoll.WaitingForInput) + manager._deferred_ack_targets[self.room] = {(target.ip, target.port): target} + + self._provoke_worker_failure(manager) + + self.assertEqual(manager._staging_outstanding[self.room], 1) + self.assertNotIn(self.room, manager._deferred_ack_targets) + + self._dispatch_abort(manager) + self.assertNotIn(self.room, manager._deferred_ack_targets) + self.assertEqual(manager._sent, []) + + manager.request_status.pop(self.room, None) + CommonKVManager.update_status(manager, self.room, KVPoll.Bootstrapping) + self._dispatch_abort(manager) + self.assertNotIn(self.room, manager._deferred_ack_targets) + self.assertEqual(manager._sent, [(self.room, target)]) + + +class TestAbortWireFormat(CustomTestCase): + def test_abort_notification_round_trip(self): + notification = AbortNotification(100, "10.0.0.1", 5000, 7) + + self.assertEqual( + AbortNotification.from_zmq(notification.to_zmq()), notification + ) + + def test_legacy_abort_without_return_address(self): + self.assertEqual( + AbortNotification.from_zmq([ABORT_TAG, b"101"]), + AbortNotification(room=101), + ) + + def test_malformed_abort_is_rejected(self): + self.assertIsNone(AbortNotification.from_zmq([ABORT_TAG, b"bad-room"])) + + def test_malformed_abort_generation_is_dropped(self): + notification = AbortNotification.from_zmq( + [ABORT_TAG, b"101", b"10.0.0.1", b"5000", b"bad-generation"] + ) + + self.assertEqual( + notification, + AbortNotification(room=101, decode_ip="10.0.0.1", decode_port=5000), + ) + + def test_abort_ack_round_trip(self): + ack = AbortAck(room=102, prefill_rank=3, generation=7) + + self.assertEqual(AbortAck.from_zmq(ack.to_zmq()), ack) + + +class TestCommonAbortAckDispatch(CustomTestCase): + def test_abort_ack_message_is_aggregated(self): + mgr = _make_manager() + generation = mgr.register_deferred_abort_room(103) + + claimed = mgr.handle_abort_ack_message(AbortAck(103, 4, generation).to_zmq()) + + self.assertTrue(claimed) + self.assertEqual(mgr._deferred_abort_ack_tracker[103].prefill_ranks, {4}) + + def test_non_ack_message_is_not_claimed(self): + mgr = _make_manager() + + self.assertFalse(mgr.handle_abort_ack_message([b"STATUS", b"103", b"4"])) + + def test_malformed_abort_ack_is_ignored(self): + mgr = _make_manager() + mgr.register_deferred_abort_room(103) + + self.assertTrue( + mgr.handle_abort_ack_message([ABORT_ACK_TAG, b"bad-room", b"4", b"1"]) + ) + self.assertFalse(mgr.is_abort_release_safe(103, required_acks=1)) + + def test_generationless_abort_ack_warns_and_is_ignored(self): + mgr = _make_manager() + mgr.register_deferred_abort_room(103) + + with patch( + "sglang.srt.disaggregation.common.conn.logger.warning_once" + ) as warning: + claimed = mgr.handle_abort_ack_message([ABORT_ACK_TAG, b"103", b"4"]) + + self.assertTrue(claimed) + self.assertFalse(mgr.is_abort_release_safe(103, required_acks=1)) + warning.assert_called_once() + + def test_stale_generation_ack_is_ignored_after_room_reuse(self): + mgr = _make_manager() + mgr.register_deferred_abort_room(103) + mgr.clear_deferred_abort_state(103) + mgr.register_deferred_abort_room(103) + + claimed = mgr.handle_abort_ack_message([ABORT_ACK_TAG, b"103", b"4", b"1"]) + + self.assertTrue(claimed) + self.assertFalse(mgr.is_abort_release_safe(103, required_acks=1)) + + claimed = mgr.handle_abort_ack_message([ABORT_ACK_TAG, b"103", b"4", b"2"]) + + self.assertTrue(claimed) + self.assertTrue(mgr.is_abort_release_safe(103, required_acks=1)) + + def _make_notifying_receiver(self, mgr, room, init_time): + receiver = _TestReceiver.__new__(_TestReceiver) + receiver.kv_mgr = mgr + receiver.bootstrap_room = room + receiver.bootstrap_infos = [{"rank_ip": "10.0.0.2", "rank_port": 6000}] + receiver.init_time = init_time + receiver.abort_notified = False + receiver.conclude_state = None + receiver._abort_generation = None + mgr.local_ip = "10.0.0.1" + mgr.rank_port = 5000 + mgr.failure_lock = threading.Lock() + mgr.failure_records = {} + mgr.request_status = {room: KVPoll.WaitingForInput} + sent = [] + + class _Sock: + def send_multipart(self, frames): + sent.append( + ( + room in mgr._deferred_abort_ack_tracker, + AbortNotification.from_zmq(frames), + ) + ) + + receiver._connect_to_bootstrap_server = lambda info: ( + _Sock(), + threading.Lock(), + ) + return receiver, sent + + def test_abort_arms_tracker_before_sending(self): + mgr = _make_manager() + receiver, sent = self._make_notifying_receiver(mgr, 104, init_time=1.0) + + receiver.abort() + + [(armed, notification)] = sent + self.assertTrue(armed) + self.assertEqual( + notification.generation, + mgr._deferred_abort_ack_tracker[104].generation, + ) + + def test_waiting_timeout_abort_is_armed(self): + mgr = _make_manager() + mgr.waiting_timeout = 0 + receiver, sent = self._make_notifying_receiver(mgr, 105, init_time=1.0) + receiver.invalidate_cached_bootstrap_infos = lambda: None + + self.assertEqual(receiver._check_waiting_timeout(), KVPoll.Failed) + + [(armed, notification)] = sent + self.assertTrue(armed) + self.assertIsNotNone(notification.generation) + self.assertTrue(receiver.abort_notified) + + def test_abort_before_metadata_skips_tracker(self): + mgr = _make_manager() + receiver, sent = self._make_notifying_receiver(mgr, 106, init_time=None) + + receiver.abort() + + [(armed, notification)] = sent + self.assertFalse(armed) + self.assertIsNone(notification.generation) + + def test_abort_skips_tracker_when_disabled(self): + mgr = _make_manager() + mgr.enable_deferred_decode_kv_release = False + receiver, sent = self._make_notifying_receiver(mgr, 107, init_time=1.0) + + receiver.abort() + + self.assertNotIn(107, mgr._deferred_abort_ack_tracker) + self.assertIsNone(sent[0][1].generation) + + +class TestDeferredAckTargets(CustomTestCase): + def test_ack_held_until_outstanding_drains(self): + mgr = _make_prefill_manager() + target = AckTarget("10.0.0.1", 5000, 9) + mgr.register_deferred_ack_target(7, target) + + mgr._staging_outstanding[7] = 1 + mgr._maybe_ack_drained_abort(7) + self.assertEqual(mgr._sent, []) + + mgr._staging_outstanding[7] = 0 + mgr._maybe_ack_drained_abort(7) + self.assertEqual(mgr._sent, [(7, target)]) + + def test_ack_fires_at_most_once(self): + mgr = _make_prefill_manager() + target = AckTarget("10.0.0.2", 5001, 10) + mgr.register_deferred_ack_target(8, target) + + mgr._maybe_ack_drained_abort(8) + mgr._maybe_ack_drained_abort(8) + + self.assertEqual(mgr._sent, [(8, target)]) + self.assertNotIn(8, mgr._deferred_ack_targets) + + def test_unregistered_room_is_noop(self): + mgr = _make_prefill_manager() + + mgr._maybe_ack_drained_abort(999) + + self.assertEqual(mgr._sent, []) + + def test_sender_clear_keeps_target_until_outstanding_transfer_drains(self): + mgr = _make_prefill_manager() + mgr.request_status = {7: KVPoll.Failed} + target = AckTarget("10.0.0.1", 5000, 9) + mgr.register_deferred_ack_target(7, target) + mgr._staging_outstanding[7] = 1 + sender = _TestSender.__new__(_TestSender) + sender.kv_mgr = mgr + sender.bootstrap_room = 7 + + sender.clear() + + self.assertEqual( + mgr._deferred_ack_targets[7], {(target.ip, target.port): target} + ) + mgr._staging_outstanding[7] = 0 + mgr._maybe_ack_drained_abort(7) + self.assertEqual(mgr._sent, [(7, target)]) + + def test_sender_clear_discards_target_without_outstanding_transfer(self): + mgr = _make_prefill_manager() + mgr.request_status = {7: KVPoll.Failed} + mgr.register_deferred_ack_target(7, AckTarget("10.0.0.1", 5000, 9)) + sender = _TestSender.__new__(_TestSender) + sender.kv_mgr = mgr + sender.bootstrap_room = 7 + + sender.clear() + + self.assertNotIn(7, mgr._deferred_ack_targets) + + def test_drain_ack_fans_out_to_every_aborting_decode_rank(self): + # Prefill TP < decode TP: two decode ranks share the room and each + # sends its own ABORT. Neither may overwrite the other's target. + mgr = _make_prefill_manager() + rank0 = AckTarget("10.0.0.1", 5000, 3) + rank1 = AckTarget("10.0.0.2", 5001, 8) + mgr._staging_outstanding[7] = 1 + mgr.register_deferred_ack_target(7, rank0) + mgr.register_deferred_ack_target(7, rank1) + + mgr._maybe_ack_drained_abort(7) + self.assertEqual(mgr._sent, []) + + mgr._staging_outstanding[7] = 0 + mgr._maybe_ack_drained_abort(7) + self.assertEqual(sorted(mgr._sent), [(7, rank0), (7, rank1)]) + mgr._maybe_ack_drained_abort(7) + self.assertEqual(len(mgr._sent), 2) + + def test_repeated_abort_from_one_decode_rank_keeps_latest_generation(self): + mgr = _make_prefill_manager() + mgr._staging_outstanding[7] = 1 + mgr.register_deferred_ack_target(7, AckTarget("10.0.0.1", 5000, 3)) + latest = AckTarget("10.0.0.1", 5000, 4) + mgr.register_deferred_ack_target(7, latest) + + mgr._staging_outstanding[7] = 0 + mgr._maybe_ack_drained_abort(7) + + self.assertEqual(mgr._sent, [(7, latest)]) + + def test_prefill_unique_rank_formula(self): + mgr = CommonKVManager.__new__(CommonKVManager) + mgr.attn_tp_rank, mgr.pp_size, mgr.attn_cp_size = 2, 3, 4 + mgr.pp_rank, mgr.attn_cp_rank = 1, 3 + + self.assertEqual(mgr._prefill_unique_rank(), 31) + + class TestAbortAckAggregation(CustomTestCase): def test_release_safe_only_after_all_required_ranks_ack(self): mgr = _make_manager() room = 100 - mgr.register_deferred_abort_room(room) + generation = mgr.register_deferred_abort_room(room) self.assertFalse(mgr.is_abort_release_safe(room, required_acks=2)) - mgr.note_abort_ack(room, 0) + mgr.note_abort_ack(room, 0, generation) self.assertFalse(mgr.is_abort_release_safe(room, required_acks=2)) - mgr.note_abort_ack(room, 1) + mgr.note_abort_ack(room, 1, generation) self.assertTrue(mgr.is_abort_release_safe(room, required_acks=2)) def test_duplicate_rank_ack_does_not_over_count(self): mgr = _make_manager() room = 101 - mgr.register_deferred_abort_room(room) - mgr.note_abort_ack(room, 0) - mgr.note_abort_ack(room, 0) # same rank twice + generation = mgr.register_deferred_abort_room(room) + mgr.note_abort_ack(room, 0, generation) + mgr.note_abort_ack(room, 0, generation) # same rank twice # Two acks arrived but from one rank: not safe for a 2-rank prefill. self.assertFalse(mgr.is_abort_release_safe(room, required_acks=2)) - def test_single_rank_fast_path(self): - mgr = _make_manager() - room = 102 - mgr.register_deferred_abort_room(room) - mgr.note_abort_ack(room, 0) - self.assertTrue(mgr.is_abort_release_safe(room, required_acks=1)) - def test_clear_deferred_abort_state(self): mgr = _make_manager() room = 103 - mgr.register_deferred_abort_room(room) - mgr.note_abort_ack(room, 0) + generation = mgr.register_deferred_abort_room(room) + mgr.note_abort_ack(room, 0, generation) mgr.clear_deferred_abort_state(room) self.assertNotIn(room, mgr._deferred_abort_ack_tracker) self.assertFalse(mgr.is_abort_release_safe(room, required_acks=1)) @@ -77,41 +530,16 @@ def test_ack_before_register_is_dropped(self): # would otherwise pollute a later request reusing the same room). mgr = _make_manager() room = 104 - mgr.note_abort_ack(room, 0) # no register yet + mgr.note_abort_ack(room, 0, 1) # no register yet self.assertNotIn(room, mgr._deferred_abort_ack_tracker) self.assertFalse(mgr.is_abort_release_safe(room, required_acks=1)) - def test_late_ack_after_release_does_not_pollute_reused_room(self): - # Regression for bootstrap_room reuse: req A (room R) releases, then a - # late ack from A arrives, then req B reuses room R. B must start from a - # clean slate and not inherit A's ack (which would release B early while - # its transfer is still in flight -> KV corruption). - mgr = _make_manager() - room = 105 - - # Req A: held, one of two ranks acks, then released (e.g. timed out). - mgr.register_deferred_abort_room(room) - mgr.note_abort_ack(room, 0) - mgr.clear_deferred_abort_state(room) - - # Late ack from A's other rank arrives after release -> dropped. - mgr.note_abort_ack(room, 1) - self.assertNotIn(room, mgr._deferred_abort_ack_tracker) - - # Req B reuses room R. - mgr.register_deferred_abort_room(room) - # Only B's rank-0 has acked so far; a 2-rank prefill is NOT safe yet. - mgr.note_abort_ack(room, 0) - self.assertFalse(mgr.is_abort_release_safe(room, required_acks=2)) - mgr.note_abort_ack(room, 1) - self.assertTrue(mgr.is_abort_release_safe(room, required_acks=2)) - def test_register_resets_stale_acks(self): mgr = _make_manager() room = 106 - mgr.register_deferred_abort_room(room) - mgr.note_abort_ack(room, 0) - mgr.note_abort_ack(room, 1) + generation = mgr.register_deferred_abort_room(room) + mgr.note_abort_ack(room, 0, generation) + mgr.note_abort_ack(room, 1, generation) self.assertTrue(mgr.is_abort_release_safe(room, required_acks=2)) # Re-registering (a later reuse) wipes the prior acks. mgr.register_deferred_abort_room(room) @@ -140,6 +568,7 @@ def _abort_receiver(self, mgr, init_time): recv.bootstrap_room = 500 recv.init_time = init_time recv.abort_notified = False + recv._abort_generation = None recv.bootstrap_infos = [{"rank_ip": "10.0.0.9", "rank_port": 7000}] armed_at_send = [] sock = SimpleNamespace( @@ -173,7 +602,7 @@ def test_prealloc_abort_does_not_arm(self): self.assertNotIn(500, mgr._deferred_abort_ack_tracker) def test_opted_out_backend_does_not_arm(self): - # Backends without a drain ack (mori) send the ABORT but must not arm: + # Backends without a drain ack send the ABORT but must not arm: # nothing would ever clean the tracker up. mgr = self._make_decode_manager(enabled=False) recv, armed_at_send = self._abort_receiver(mgr, init_time=123.0) @@ -247,9 +676,11 @@ def test_host_release_requires_drain_on_every_rank_even_after_timeout(self): q.enable_host_receive = True q.gloo_group = object() entries = [_make_decode_req(room, room, mgr) for room in (1, 2)] + generations = {} for entry in entries: entry.host_staged = True - mgr.register_deferred_abort_room(entry.req.bootstrap_room) + room = entry.req.bootstrap_room + generations[room] = mgr.register_deferred_abort_room(room) q._defer_release(entry) with ( patch.object(decode_mod, "discard_kv_cache_backup") as discard, @@ -259,7 +690,7 @@ def test_host_release_requires_drain_on_every_rank_even_after_timeout(self): ): q.resolve_deferred_releases() discard.assert_not_called() - mgr.note_abort_ack(2, 0) + mgr.note_abort_ack(2, 0, generations[2]) reduce.side_effect = lambda ready, **_: ready.zero_() q.resolve_deferred_releases() discard.assert_not_called() @@ -267,7 +698,7 @@ def test_host_release_requires_drain_on_every_rank_even_after_timeout(self): q.resolve_deferred_releases() discard.assert_called_once_with(entries[1].req, q.tree_cache, "host_pool") self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [2]) - mgr.note_abort_ack(1, 0) + mgr.note_abort_ack(1, 0, generations[1]) q.resolve_deferred_releases() self.assertEqual(discard.call_count, 2) self.assertEqual(q._deferred_releases, []) @@ -286,7 +717,7 @@ def test_holds_until_drained_then_releases(self): dreq = _make_decode_req(room, idx, mgr, n_prefill_ranks=2) # In production the receiver arms the room just before it sends the # ABORT (_send_abort_notification), before the scheduler defers here. - mgr.register_deferred_abort_room(room) + generation = mgr.register_deferred_abort_room(room) q._defer_release(dreq) with patch.object(decode_mod, "release_kv_cache") as rel: @@ -296,13 +727,13 @@ def test_holds_until_drained_then_releases(self): self.assertEqual(len(q._deferred_releases), 1) # One of two ranks acked -> still held. - mgr.note_abort_ack(room, 0) + mgr.note_abort_ack(room, 0, generation) q.resolve_deferred_releases() rel.assert_not_called() self.assertEqual(len(q._deferred_releases), 1) # Both ranks acked -> released exactly once. - mgr.note_abort_ack(room, 1) + mgr.note_abort_ack(room, 1, generation) q.resolve_deferred_releases() rel.assert_called_once_with(dreq.req, q.tree_cache, is_insert=False) @@ -483,14 +914,15 @@ def test_enabled_by_default(self): self.assertTrue(envs.SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE.get()) def test_backends_that_ack_opt_in(self): - # Ascend inherits Mooncake's threads, so it opts in too. + # Ascend inherits Mooncake's threads, so it opts in too. Mori's opt-in + # is asserted in test_mori_deferred_kv_release.py, which stubs mori. for cls in (MooncakeKVManager, NixlKVManager): with self.subTest(backend=cls.__name__): self.assertTrue(cls.supports_deferred_decode_kv_release) def test_backends_without_a_drain_ack_stay_opted_out(self): - # Inheriting CommonKVManager is not enough: mori marks the room Failed - # without acking. + # Inheriting CommonKVManager is not enough: a backend must send the + # drain ack itself before it may opt in. self.assertFalse(BaseKVManager.supports_deferred_decode_kv_release) self.assertFalse(CommonKVManager.supports_deferred_decode_kv_release) diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 3663a2f689e9..9db56f33468a 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -10,7 +10,7 @@ import torch.distributed as dist from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll, StateType -from sglang.srt.disaggregation.common.conn import CommonKVManager +from sglang.srt.disaggregation.common.conn import AckTarget, CommonKVManager from sglang.srt.disaggregation.common.staging_buffer import ( StagingAllocator, ) @@ -333,20 +333,21 @@ def test_sender_clear_keeps_abort_ack_until_writes_drain(self): manager = object.__new__(MooncakeKVManager) sender = object.__new__(MooncakeKVSender) sender.kv_mgr, sender.bootstrap_room = manager, 42 + target = AckTarget("127.0.0.1", 1234, 7) for outstanding in (0, 1): with self.subTest(outstanding=outstanding): manager.request_status = {42: KVPoll.Failed} - manager._staging_outstanding = {42: outstanding} - manager._deferred_ack_targets = {42: ("127.0.0.1", 1234)} - manager._deferred_ack_fanout_snapshots = {} + manager.req_to_decode_prefix_len = {} manager.transfer_infos = {} + manager._staging_outstanding = {42: outstanding} + manager._deferred_ack_targets = {42: {(target.ip, target.port): target}} with patch.object(manager, "_send_abort_ack") as ack: sender.clear() if outstanding: ack.assert_not_called() manager._staging_outstanding[42] = 0 manager._maybe_ack_drained_abort(42) - ack.assert_called_once_with("127.0.0.1", 1234, 42) + ack.assert_called_once_with(42, target) manager._maybe_ack_drained_abort(42) ack.assert_called_once() diff --git a/test/registered/unit/disaggregation/test_mooncake_deferred_kv_release.py b/test/registered/unit/disaggregation/test_mooncake_deferred_kv_release.py new file mode 100644 index 000000000000..705156ec576c --- /dev/null +++ b/test/registered/unit/disaggregation/test_mooncake_deferred_kv_release.py @@ -0,0 +1,146 @@ +"""Deferred decode-side KV release on the Mooncake backend. + +The bootstrap thread records the ACK target while a transfer is outstanding. +Only the transfer worker can send that ACK after its write has drained. +""" + +import unittest +from collections import defaultdict +from unittest.mock import MagicMock, patch + +import numpy as np +from test_deferred_decode_kv_release import ( + DeferredAbortNotificationScenarios, + TaggedAbortNotificationScenarios, + WorkerFailureAbortScenarios, +) + +from sglang.srt.disaggregation.base.conn import KVPoll +from sglang.srt.disaggregation.common.conn import ( + ABORT_ACK_TAG, + ABORT_TAG, + AbortNotification, +) +from sglang.srt.disaggregation.common.utils import TransferKVChunk +from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +ROOM = 106 +DECODE_IP = "10.0.0.2" +DECODE_PORT = 6000 + + +def _manager( + *, enabled: bool = True, status: KVPoll | None = KVPoll.WaitingForInput +) -> MooncakeKVManager: + manager = MooncakeKVManager.__new__(MooncakeKVManager) + manager.enable_deferred_decode_kv_release = enabled + manager.request_status = {} if status is None else {ROOM: status} + manager._deferred_ack_targets = {} + manager._deferred_ack_poisoned_rooms = set() + manager._staging_outstanding = {} + manager.check_status = lambda room: manager.request_status[room] + manager.update_status = lambda room, value: manager.request_status.__setitem__( + room, value + ) + manager._sent = [] + manager._send_abort_ack = lambda *args: manager._sent.append(args) + manager._send_multipart_locked = MagicMock() + return manager + + +def _chunk() -> TransferKVChunk: + return TransferKVChunk( + room=ROOM, + prefill_kv_indices=np.array([1], dtype=np.int32), + index_slice=slice(0, 1), + is_last_chunk=False, + chunk_id=0, + prefill_aux_index=None, + state_indices=None, + ) + + +class TestMooncakeAbortNotification( + DeferredAbortNotificationScenarios, + TaggedAbortNotificationScenarios, + WorkerFailureAbortScenarios, + CustomTestCase, +): + room = ROOM + decode_ip = DECODE_IP + decode_port = DECODE_PORT + + def _make_abort_manager(self, status: KVPoll | None): + return _manager(status=status) + + def _provoke_worker_failure(self, manager) -> None: + manager.enable_trace = False + manager.bootstrap_port = 8998 + manager._staging_outstanding = defaultdict(int) + manager.check_status = MagicMock(side_effect=RuntimeError("worker error")) + queue = MagicMock() + queue.get.return_value = _chunk() + + with self.assertRaisesRegex(RuntimeError, "Transfer thread failed"): + manager.transfer_worker(queue, MagicMock()) + manager.check_status = lambda room: manager.request_status[room] + + def test_generationless_abort_marks_failed_without_ack_or_warning(self): + # A decode that did not arm (prealloc abort) sends no generation; that + # is normal and must not warn about mixed versions. + manager = _manager() + message = AbortNotification(ROOM, DECODE_IP, DECODE_PORT).to_zmq() + + with patch( + "sglang.srt.disaggregation.common.conn.logger.warning_once" + ) as warning: + claimed = manager._handle_abort_notification(message) + + self.assertTrue(claimed) + self.assertEqual(manager.request_status[ROOM], KVPoll.Failed) + self.assertEqual(manager._deferred_ack_targets, {}) + self.assertEqual(manager._sent, []) + warning.assert_not_called() + + def test_abort_without_valid_return_address_still_marks_room_failed(self): + messages = ( + ("missing", [ABORT_TAG, str(ROOM).encode("ascii")]), + ( + "malformed port", + [ + ABORT_TAG, + str(ROOM).encode("ascii"), + DECODE_IP.encode("ascii"), + b"bad", + ], + ), + ) + + for name, message in messages: + with self.subTest(name=name): + manager = _manager() + + claimed = manager._handle_abort_notification(message) + + self.assertTrue(claimed) + self.assertEqual(manager.request_status[ROOM], KVPoll.Failed) + self.assertEqual(manager._deferred_ack_targets, {}) + self.assertEqual(manager._sent, []) + + def test_feature_off_sends_legacy_ack(self): + manager = _manager(enabled=False) + + claimed = manager._handle_abort_notification(self._abort_message()) + + self.assertTrue(claimed) + self.assertEqual(manager.request_status[ROOM], KVPoll.Failed) + payload = manager._send_multipart_locked.call_args.args[1] + self.assertEqual(payload, [ABORT_ACK_TAG, str(ROOM).encode("ascii")]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/disaggregation/test_mori_deferred_kv_release.py b/test/registered/unit/disaggregation/test_mori_deferred_kv_release.py new file mode 100644 index 000000000000..1c20e45aded9 --- /dev/null +++ b/test/registered/unit/disaggregation/test_mori_deferred_kv_release.py @@ -0,0 +1,376 @@ +"""Deferred decode-side KV release on the Mori backend.""" + +import importlib +import sys +import threading +import types +import unittest +from collections import defaultdict +from enum import IntEnum +from queue import Queue +from unittest.mock import MagicMock, patch + +import numpy as np + + +def _install_mori_stubs() -> None: + try: + importlib.import_module("mori.cpp") + importlib.import_module("mori.io") + return + except (ImportError, OSError): + pass + + mori = types.ModuleType("mori") + mori.__path__ = [] + mori_cpp = types.ModuleType("mori.cpp") + mori_io = types.ModuleType("mori.io") + + class FakeStatusCode(IntEnum): + SUCCESS = 0 + IN_PROGRESS = 1 + FAILED = 2 + + mori_cpp.TransferStatus = object + for name in ( + "BackendType", + "EngineDesc", + "IOEngine", + "IOEngineConfig", + "MemoryDesc", + "MemoryLocationType", + "PollCqMode", + "RdmaBackendConfig", + ): + setattr(mori_io, name, object) + mori_io.StatusCode = FakeStatusCode + mori.cpp = mori_cpp + mori.io = mori_io + sys.modules["mori"] = mori + sys.modules["mori.cpp"] = mori_cpp + sys.modules["mori.io"] = mori_io + + +_install_mori_stubs() + +from test_deferred_decode_kv_release import ( + ABORT_GENERATION, + DeferredAbortNotificationScenarios, +) + +from sglang.srt.disaggregation.base.conn import KVPoll +from sglang.srt.disaggregation.common.conn import ( + ABORT_TAG, + AbortNotification, + AckTarget, +) +from sglang.srt.disaggregation.common.utils import TransferKVChunk +from sglang.srt.disaggregation.mori.conn import ( + MoriKVManager, + StatusCode, + _MoriTransferSubmissionError, + _SubmissionLocal, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _manager(enabled: bool = True) -> MoriKVManager: + manager = MoriKVManager.__new__(MoriKVManager) + manager.enable_deferred_decode_kv_release = enabled + manager.request_status = {} + manager.transfer_lock = threading.Lock() + manager._abort_ack_lock = threading.Lock() + manager._deferred_ack_targets = {} + manager._deferred_ack_poisoned_rooms = set() + manager._deferred_abort_ack_tracker = {} + manager._staging_outstanding = defaultdict(int) + manager._drain_queue = Queue(maxsize=1) + manager._submission_local = _SubmissionLocal() + manager.req_to_decode_prefix_len = {} + manager.transfer_infos = {} + manager.failure_lock = threading.Lock() + manager.failure_records = {} + manager._sent = [] + manager._send_abort_ack = lambda *args: manager._sent.append(args) + manager._wait_poll_ms = 0 + manager._transfer_timeout_ms = 0 + return manager + + +def _abort_message(room: int = 11) -> list[bytes]: + return AbortNotification(room, "10.0.0.3", 6000, ABORT_GENERATION).to_zmq() + + +def _chunk(room: int = 11) -> TransferKVChunk: + return TransferKVChunk( + room=room, + prefill_kv_indices=np.array([1], dtype=np.int32), + index_slice=slice(0, 1), + is_last_chunk=False, + prefill_aux_index=None, + state_indices=None, + ) + + +class TestMoriAbortAck(DeferredAbortNotificationScenarios, CustomTestCase): + room = 11 + decode_ip = "10.0.0.3" + decode_port = 6000 + + def _make_abort_manager(self, status: KVPoll | None): + manager = _manager() + if status is not None: + manager.request_status[self.room] = status + return manager + + def _dispatch_abort(self, manager) -> None: + manager._handle_abort_message(self._abort_message()) + + def _start_test_transfer(self, manager) -> None: + manager._mark_transfer_started(_chunk(self.room)) + + def _drain_test_transfer(self, manager) -> None: + manager._mark_transfer_quiescent(_chunk(self.room)) + + def test_abort_without_valid_return_address_marks_failed_without_ack(self): + messages = ( + ("missing", [ABORT_TAG, b"11"]), + ("malformed port", [ABORT_TAG, b"11", b"10.0.0.3", b"bad"]), + ) + + for name, message in messages: + with self.subTest(name=name): + manager = _manager() + manager.request_status[11] = KVPoll.WaitingForInput + + manager._handle_abort_message(message) + + self.assertEqual(manager.request_status[11], KVPoll.Failed) + self.assertEqual(manager._sent, []) + self.assertEqual(manager._deferred_ack_targets, {}) + + def test_feature_off_preserves_abort_without_ack(self): + manager = _manager(enabled=False) + manager.request_status[11] = KVPoll.WaitingForInput + + manager._handle_abort_message(_abort_message()) + + self.assertEqual(manager.request_status[11], KVPoll.Failed) + self.assertEqual(manager._sent, []) + self.assertEqual(manager._deferred_ack_targets, {}) + + def test_abort_after_submit_waits_for_terminal_status(self): + manager = _manager() + manager.request_status[11] = KVPoll.Transferring + status = MagicMock() + status.InProgress.return_value = False + manager._submit_kv_transfer = MagicMock(return_value=[status]) + wait_results = iter((StatusCode.IN_PROGRESS, StatusCode.SUCCESS)) + + def wait_all(*args, **kwargs): + result = next(wait_results) + if result == StatusCode.IN_PROGRESS: + manager._handle_abort_message(_abort_message()) + self.assertEqual(manager._sent, []) + return result + + manager.engine = MagicMock() + manager.engine.wait_all.side_effect = wait_all + chunk = _chunk() + manager._mark_transfer_started(chunk) + + is_quiescent = manager._process_transfer_chunk(chunk) + + self.assertTrue(is_quiescent) + self.assertEqual(manager.engine.wait_all.call_count, 2) + self.assertEqual(manager._sent, []) + manager._mark_transfer_quiescent(chunk) + self.assertEqual( + manager._sent, + [(11, AckTarget("10.0.0.3", 6000, ABORT_GENERATION))], + ) + + def test_sla_failure_acks_after_status_drains(self): + manager = _manager() + manager.request_status[11] = KVPoll.Transferring + manager._transfer_timeout_ms = 1 + manager.engine = MagicMock() + manager.engine.wait_all.return_value = StatusCode.IN_PROGRESS + status = MagicMock() + status.InProgress.return_value = True + manager._submit_kv_transfer = MagicMock(return_value=[status]) + manager.conclude_failure = MagicMock() + chunk = _chunk() + manager._mark_transfer_started(chunk) + + with patch( + "sglang.srt.disaggregation.mori.conn.time.perf_counter", + side_effect=(0.0, 0.002), + ): + is_quiescent = manager._process_transfer_chunk(chunk) + + self.assertFalse(is_quiescent) + manager.engine.wait_all.assert_called_once() + self.assertEqual(manager._sent, []) + self.assertEqual(manager._staging_outstanding[11], 1) + # Concluded at detection; only the ACK waits for the drain. + manager.conclude_failure.assert_called_once_with( + bootstrap_room=11, failure_reason="KV transfer exceeded SLA 1ms" + ) + manager._handle_abort_message(_abort_message()) + queued_chunk, queued_statuses = manager._drain_queue.get_nowait() + manager._drain_transfer_statuses(queued_chunk, queued_statuses) + + status.Wait.assert_called_once_with() + self.assertEqual( + manager._sent, + [(11, AckTarget("10.0.0.3", 6000, ABORT_GENERATION))], + ) + self.assertNotIn(11, manager._staging_outstanding) + manager.conclude_failure.assert_called_once() + + def test_room_cleared_after_drain_handoff_keeps_single_owner(self): + manager = _manager() + manager.request_status[11] = KVPoll.Transferring + manager._transfer_timeout_ms = 1 + status = MagicMock() + status.InProgress.return_value = True + manager._submit_kv_transfer = MagicMock(return_value=[status]) + + def wait_all(*args, **kwargs): + manager.request_status.pop(11) + return StatusCode.IN_PROGRESS + + manager.engine = MagicMock() + manager.engine.wait_all.side_effect = wait_all + manager.conclude_failure = MagicMock() + queue = MagicMock() + queue.get.side_effect = (_chunk(), KeyboardInterrupt()) + + with patch( + "sglang.srt.disaggregation.mori.conn.time.perf_counter", + side_effect=(0.0, 0.002), + ): + with self.assertRaises(KeyboardInterrupt): + manager._transfer_worker(queue) + + self.assertEqual(manager._staging_outstanding[11], 1) + self.assertEqual(manager._drain_queue.qsize(), 1) + manager.conclude_failure.assert_called_once() + + def test_sla_failure_keeps_legacy_early_return_when_disabled(self): + manager = _manager(enabled=False) + manager._transfer_timeout_ms = 1 + manager.engine = MagicMock() + manager.engine.wait_all.side_effect = ( + StatusCode.IN_PROGRESS, + StatusCode.SUCCESS, + ) + + with patch( + "sglang.srt.disaggregation.mori.conn.time.perf_counter", + side_effect=(0.0, 0.002), + ): + failure = manager._wait_transfer_completion([object()]) + + self.assertEqual(failure, "KV transfer exceeded SLA 1ms") + self.assertEqual(manager.engine.wait_all.call_count, 1) + + def test_wait_event_failure_releases_outstanding_count(self): + manager = _manager() + manager.request_status[11] = KVPoll.Transferring + manager.conclude_failure = MagicMock() + chunk = _chunk() + chunk.wait_event = MagicMock() + chunk.wait_event.synchronize.side_effect = RuntimeError("event failed") + queue = MagicMock() + queue.get.side_effect = (chunk, KeyboardInterrupt()) + + with patch("sglang.srt.disaggregation.mori.conn.logger.exception"): + with self.assertRaises(KeyboardInterrupt): + manager._transfer_worker(queue) + + self.assertNotIn(11, manager._staging_outstanding) + manager.conclude_failure.assert_called_once() + + def test_wait_all_exception_keeps_status_for_draining(self): + manager = _manager() + manager.request_status[11] = KVPoll.Transferring + status = MagicMock() + manager._submit_kv_transfer = MagicMock(return_value=[status]) + manager.engine = MagicMock() + manager.engine.wait_all.side_effect = RuntimeError("wait failed") + manager.conclude_failure = MagicMock() + + is_quiescent = manager._process_transfer_chunk(_chunk()) + + self.assertFalse(is_quiescent) + queued_chunk, queued_statuses = manager._drain_queue.get_nowait() + self.assertEqual(queued_chunk.room, 11) + self.assertEqual(queued_statuses, [status]) + manager.conclude_failure.assert_called_once_with( + bootstrap_room=11, + failure_reason="Transfer completion failed: RuntimeError('wait failed')", + ) + + def test_partial_submission_failure_transfers_ownership_to_drainer(self): + manager = _manager() + manager.request_status[11] = KVPoll.Transferring + status = MagicMock() + + def submit(*args, **kwargs): + manager._record_submitted_statuses([status]) + raise RuntimeError("later target failed") + + manager._submit_kv_transfer = submit + manager.conclude_failure = MagicMock() + chunk = _chunk() + manager._mark_transfer_started(chunk) + + with self.assertRaises(_MoriTransferSubmissionError) as raised: + manager._process_transfer_chunk(chunk) + manager._handle_submission_failure(chunk, raised.exception) + + queued_chunk, queued_statuses = manager._drain_queue.get_nowait() + self.assertIs(queued_chunk, chunk) + self.assertEqual(queued_statuses, [status]) + self.assertEqual(manager._staging_outstanding[11], 1) + manager.conclude_failure.assert_called_once_with( + bootstrap_room=11, + failure_reason="Transfer submission failed: later target failed", + ) + + def test_drain_worker_retries_same_item_after_exception(self): + manager = _manager() + manager._wait_poll_ms = 0 + chunk = _chunk() + status = MagicMock() + manager._drain_queue = MagicMock() + manager._drain_queue.get.side_effect = ( + (chunk, [status]), + KeyboardInterrupt(), + ) + manager._drain_transfer_statuses = MagicMock( + side_effect=(RuntimeError("transient wait failure"), None) + ) + + with patch("sglang.srt.disaggregation.mori.conn.logger.exception") as log: + with patch("sglang.srt.disaggregation.mori.conn.time.sleep"): + with self.assertRaises(KeyboardInterrupt): + manager._drain_worker() + + self.assertEqual(manager._drain_transfer_statuses.call_count, 2) + log.assert_called_once() + + +class TestMoriBackendOptIn(CustomTestCase): + def test_mori_opts_into_deferred_release(self): + # Mori sends the drain ack itself, so the default-on env var applies. + self.assertTrue(MoriKVManager.supports_deferred_decode_kv_release) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py b/test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py index 6c66d022dbf1..9fb502208d5b 100644 --- a/test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py +++ b/test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py @@ -1,325 +1,182 @@ """Deferred decode-side KV release on the NIXL backend. -When a decode request is aborted while its prefill->decode transfer may still be -in flight, the decode holds its KV pages until every prefill rank acks that its -transfer drained. NIXL transfers are asynchronous (agent.transfer() posts, the -worker polls check_xfer_state), so the ack must come from the transfer worker -after its DONE barrier -- never from the bootstrap thread for an active room. +NIXL posts transfers asynchronously and the worker polls their handles. The +bootstrap thread records the ACK target while a transfer is outstanding; only +the worker can send that ACK after every handle reaches DONE. """ import unittest from collections import defaultdict from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch + +import numpy as np +from test_deferred_decode_kv_release import ( + ABORT_GENERATION, + DeferredAbortNotificationScenarios, + TaggedAbortNotificationScenarios, + WorkerFailureAbortScenarios, +) from sglang.srt.disaggregation.base.conn import KVPoll -from sglang.srt.disaggregation.common.conn import CommonKVManager +from sglang.srt.disaggregation.common.conn import ( + ABORT_TAG, + AckTarget, +) +from sglang.srt.disaggregation.common.utils import TransferKVChunk from sglang.srt.disaggregation.nixl.conn import NixlKVManager from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=11, suite="base-a-test-cpu") - -def _prefill_mgr(cls=CommonKVManager, enabled=True): - """Bare manager carrying only the prefill-side deferred-ack state.""" - mgr = cls.__new__(cls) - mgr.enable_deferred_decode_kv_release = enabled - mgr._deferred_ack_targets = {} - mgr._deferred_ack_fanout_snapshots = {} - mgr._staging_outstanding = {} - mgr.request_status = {} - mgr.transfer_infos = {} - mgr._sent = [] - # Capture acks instead of opening a socket. - mgr._send_abort_ack = lambda ip, port, room: mgr._sent.append((ip, port, room)) - return mgr - - -class TestDeferredAckTargets(CustomTestCase): - def test_ack_held_until_outstanding_drains(self): - mgr = _prefill_mgr() - mgr.register_deferred_ack_target(7, "10.0.0.1", 5000) - - mgr._staging_outstanding[7] = 1 - mgr._maybe_ack_drained_abort(7) - self.assertEqual(mgr._sent, []) # still writing -> no ack - - mgr._staging_outstanding[7] = 0 - mgr._maybe_ack_drained_abort(7) - self.assertEqual(mgr._sent, [("10.0.0.1", 5000, 7)]) - - def test_ack_fires_at_most_once(self): - mgr = _prefill_mgr() - mgr.register_deferred_ack_target(8, "10.0.0.2", 5001) - mgr._maybe_ack_drained_abort(8) - mgr._maybe_ack_drained_abort(8) - self.assertEqual(len(mgr._sent), 1) - self.assertNotIn(8, mgr._deferred_ack_targets) - - def test_unregistered_room_is_noop(self): - mgr = _prefill_mgr() - mgr._maybe_ack_drained_abort(999) - self.assertEqual(mgr._sent, []) - - def test_drain_ack_fans_out_to_every_room_peer(self): - """With prefill TP < decode TP the room has several decode peers but the - registry keeps only the last ABORT sender; the drain ack must reach every - peer (dummy pairings included) or the others hold until the timeout.""" - mgr = _prefill_mgr() - mgr.transfer_infos[7] = { - "sess0": SimpleNamespace(endpoint="10.0.0.1", dst_port=5000), - "sess1": SimpleNamespace(endpoint="10.0.0.2", dst_port=5001), - # Dummy pairing: the decode rank still counts this prefill's ack. - "sess2": SimpleNamespace(endpoint="10.0.0.3", dst_port=5002), - } - # Rank 1 registered last and overwrote rank 0's registration. - mgr.register_deferred_ack_target(7, "10.0.0.2", 5001) - - mgr._maybe_ack_drained_abort(7) - self.assertEqual( - sorted(mgr._sent), - [ - ("10.0.0.1", 5000, 7), - ("10.0.0.2", 5001, 7), - ("10.0.0.3", 5002, 7), - ], +ROOM = 11 +DECODE_IP = "10.0.0.3" +DECODE_PORT = 6000 + + +def _manager( + *, enabled: bool = True, status: KVPoll | None = KVPoll.WaitingForInput +) -> NixlKVManager: + manager = NixlKVManager.__new__(NixlKVManager) + manager.enable_deferred_decode_kv_release = enabled + manager._deferred_ack_targets = {} + manager._deferred_ack_poisoned_rooms = set() + manager._staging_outstanding = {} + manager.request_status = {} if status is None else {ROOM: status} + manager._sent = [] + manager._send_abort_ack = lambda *args: manager._sent.append(args) + manager.record_failure = MagicMock() + manager.update_status = MagicMock( + side_effect=lambda room, value: manager.request_status.__setitem__(room, value) + ) + manager.check_status = lambda room: manager.request_status[room] + return manager + + +def _chunk() -> TransferKVChunk: + return TransferKVChunk( + room=ROOM, + prefill_kv_indices=np.array([1], dtype=np.int32), + index_slice=slice(0, 1), + is_last_chunk=False, + chunk_id=0, + prefill_aux_index=None, + state_indices=None, + ) + + +class TestNixlAbortNotification( + DeferredAbortNotificationScenarios, + TaggedAbortNotificationScenarios, + WorkerFailureAbortScenarios, + CustomTestCase, +): + room = ROOM + decode_ip = DECODE_IP + decode_port = DECODE_PORT + + def _make_abort_manager(self, status: KVPoll | None): + return _manager(status=status) + + def _provoke_worker_failure(self, manager) -> None: + manager.enable_staging = False + manager._staging_ctx = None + manager._staging_outstanding = defaultdict(int) + manager.transfer_infos = {ROOM: {}} + manager.exceptions = {} + # A handle that never settles may still be writing, so the chunk stays + # counted and the room is poisoned instead of notified. + manager.conclude_failure = MagicMock() + manager._await_handles = MagicMock(return_value=(False, True)) + manager.check_status = MagicMock(side_effect=RuntimeError("worker error")) + queue = SimpleNamespace( + get=MagicMock(side_effect=(_chunk(), KeyboardInterrupt())) ) - # pop() semantics survive the fan-out: a second drain acks nobody. - mgr._maybe_ack_drained_abort(7) - self.assertEqual(len(mgr._sent), 3) - - def test_drain_ack_fanout_survives_mid_flight_teardown(self): - """The sender's clear() can pop transfer_infos while a chunk is still in - flight; the peers snapshotted when the ABORT registered must still be - acked when the worker finally drains, or they hold until the timeout.""" - mgr = _prefill_mgr() - mgr.transfer_infos[7] = { - "sess0": SimpleNamespace(endpoint="10.0.0.1", dst_port=5000), - "sess1": SimpleNamespace(endpoint="10.0.0.2", dst_port=5001), - } - mgr._staging_outstanding[7] = 1 - mgr.register_deferred_ack_target(7, "10.0.0.2", 5001) - mgr._maybe_ack_drained_abort(7) - self.assertEqual(mgr._sent, []) # still writing -> held - - # Scheduler clears the sender mid-flight; the worker drains after. - mgr.transfer_infos.clear() - mgr._staging_outstanding[7] = 0 - mgr._maybe_ack_drained_abort(7) + with patch("sglang.srt.disaggregation.nixl.conn.logger.exception"): + with self.assertRaises(KeyboardInterrupt): + manager.transfer_worker(queue) + manager.check_status = lambda room: manager.request_status[room] + + def test_worker_skip_between_failure_and_target_registration_still_acks(self): + manager = _manager() + manager._staging_outstanding[ROOM] = 1 + update_status = manager.update_status.side_effect + + def fail_then_skip(room, status): + update_status(room, status) + # Window 1: the worker skips after Failed but before target + # registration, so the bootstrap thread must observe zero and ACK. + manager._staging_outstanding.pop(room, None) + manager._maybe_ack_drained_abort(room) + + manager.update_status = MagicMock(side_effect=fail_then_skip) + + claimed = manager._handle_abort_notification(self._abort_message()) + + self.assertTrue(claimed) self.assertEqual( - sorted(mgr._sent), - [("10.0.0.1", 5000, 7), ("10.0.0.2", 5001, 7)], + manager._sent, + [(ROOM, AckTarget(DECODE_IP, DECODE_PORT, ABORT_GENERATION))], ) - self.assertNotIn(7, mgr._deferred_ack_fanout_snapshots) - - def test_drain_ack_after_teardown_falls_back_to_registered_target(self): - mgr = _prefill_mgr() - mgr.register_deferred_ack_target(9, "10.0.0.4", 5003) - mgr._maybe_ack_drained_abort(9) - self.assertEqual(mgr._sent, [("10.0.0.4", 5003, 9)]) - - def test_populated_room_without_registration_acks_nobody(self): - """A normal (non-aborted) drain must not fan out spurious acks.""" - mgr = _prefill_mgr() - mgr.transfer_infos[12] = { - "sess0": SimpleNamespace(endpoint="10.0.0.5", dst_port=5004), - } - mgr._maybe_ack_drained_abort(12) - self.assertEqual(mgr._sent, []) - - def test_prefill_unique_rank_matches_success_sync_formula(self): - mgr = CommonKVManager.__new__(CommonKVManager) - mgr.attn_tp_rank, mgr.pp_size, mgr.attn_cp_size = 2, 3, 4 - mgr.pp_rank, mgr.attn_cp_rank = 1, 3 - self.assertEqual(mgr._prefill_unique_rank(), 2 * (3 * 4) + 1 * 4 + 3) - - -class TestNixlAbortNotification(CustomTestCase): - """_handle_abort_notification is the prefill bootstrap-thread entry point.""" - - @staticmethod - def _abort_msg(room=11, ip="10.0.0.3", port=6000): - return [ - b"ABORT", - str(room).encode("ascii"), - ip.encode("ascii"), - str(port).encode("ascii"), - ] - - def _mgr(self, enabled=True, room=11, status=KVPoll.WaitingForInput): - mgr = _prefill_mgr(NixlKVManager, enabled=enabled) - if status is not None: - mgr.request_status[room] = status - mgr.record_failure = MagicMock() - mgr.update_status = MagicMock( - side_effect=lambda r, s: mgr.request_status.__setitem__(r, s) + self.assertEqual(manager._deferred_ack_targets, {}) + + def test_settled_worker_failure_uncounts_chunk_and_acks(self): + # Every handle settled, so the writes are over: decode is told via + # conclude_failure, and its ABORT, which lands while the chunk is + # still counted, must be ACKed once the worker uncounts the chunk. + manager = _manager() + manager.enable_staging = False + manager._staging_ctx = None + manager._staging_outstanding = defaultdict(int) + manager.transfer_infos = {ROOM: {}} + manager.exceptions = {} + manager.check_status = MagicMock(side_effect=RuntimeError("worker error")) + + def notify_decode(**kwargs): + manager.check_status = lambda room: manager.request_status[room] + self._dispatch_abort(manager) + + manager.conclude_failure = MagicMock(side_effect=notify_decode) + queue = SimpleNamespace( + get=MagicMock(side_effect=(_chunk(), KeyboardInterrupt())) ) - mgr.check_status = lambda r: mgr.request_status[r] - return mgr - - def test_in_flight_room_registers_target_and_does_not_ack_yet(self): - # A counted chunk holds the ack: only the worker knows when it landed. - mgr = self._mgr() - mgr._staging_outstanding[11] = 1 - self.assertTrue(mgr._handle_abort_notification(self._abort_msg())) - - self.assertEqual(mgr._deferred_ack_targets[11], ("10.0.0.3", 6000)) - self.assertEqual(mgr._sent, []) - # Marked Failed first, so no new chunk can be enqueued for the room. - self.assertEqual(mgr.request_status[11], KVPoll.Failed) - def test_quiescent_active_room_acks_without_waiting_for_a_worker_visit(self): - # Window 2: chunks already drained with none left to come, so the worker - # never revisits the room -- acking here keeps it off the timeout path. - mgr = self._mgr() - self.assertTrue(mgr._handle_abort_notification(self._abort_msg())) - - self.assertEqual(mgr._sent, [("10.0.0.3", 6000, 11)]) - self.assertEqual(mgr._deferred_ack_targets, {}) + with patch("sglang.srt.disaggregation.nixl.conn.logger.exception"): + with self.assertRaises(KeyboardInterrupt): + manager.transfer_worker(queue) - def test_worker_skip_before_registration_still_acks(self): - # Window 1: the worker can pass its skip point between the Failed flip - # and registration; the ack attempt at registration covers that. - mgr = self._mgr() - mgr._staging_outstanding[11] = 1 - - real_update = mgr.update_status.side_effect - - def failed_then_worker_skips(room, status): - real_update(room, status) - # Worker dequeues, sees Failed, uncounts, and finds no target yet. - mgr._staging_outstanding.pop(room, None) - mgr._maybe_ack_drained_abort(room) - - mgr.update_status = MagicMock(side_effect=failed_then_worker_skips) - self.assertTrue(mgr._handle_abort_notification(self._abort_msg())) - - self.assertEqual(mgr._sent, [("10.0.0.3", 6000, 11)]) - self.assertEqual(mgr._deferred_ack_targets, {}) - - def test_concluded_room_acks_immediately(self): - # Concluded and quiescent: ack straight away. - mgr = self._mgr(status=None) - mgr.check_status = lambda r: KVPoll.Success - self.assertTrue(mgr._handle_abort_notification(self._abort_msg())) - - self.assertEqual(mgr._sent, [("10.0.0.3", 6000, 11)]) - self.assertEqual(mgr._deferred_ack_targets, {}) - - def test_cleared_room_with_outstanding_chunk_does_not_ack(self): - # The ERR path abandons sibling handles that may still be writing and - # leaves the chunk counted; clear() then drops the room. Acking on - # "unknown room" alone would release decode pages under those writes. - mgr = self._mgr(status=None) # room absent == cleared/unknown - mgr._staging_outstanding[11] = 1 - self.assertTrue(mgr._handle_abort_notification(self._abort_msg())) - - self.assertEqual(mgr._sent, []) - self.assertEqual(mgr._deferred_ack_targets, {}) - - def test_feature_off_registers_nothing_and_acks_nothing(self): - mgr = self._mgr(enabled=False) - self.assertTrue(mgr._handle_abort_notification(self._abort_msg())) - - self.assertEqual(mgr._deferred_ack_targets, {}) - self.assertEqual(mgr._sent, []) - # Legacy behavior preserved: the room is still failed. - self.assertEqual(mgr.request_status[11], KVPoll.Failed) - - def test_legacy_two_frame_abort_is_tolerated(self): - # Older peers send [ABORT, room] with no return address. - mgr = self._mgr() - self.assertTrue(mgr._handle_abort_notification([b"ABORT", b"11"])) - self.assertEqual(mgr._deferred_ack_targets, {}) - self.assertEqual(mgr._sent, []) - - def test_non_abort_message_is_not_claimed(self): - mgr = self._mgr() - self.assertFalse(mgr._handle_abort_notification([b"STAGING_REQ", b"11"])) - - -class _StopWorker(Exception): - pass - - -class _OneChunkQueue: - """Feeds the worker a single chunk, then unblocks it out of its loop.""" - - def __init__(self, chunk): - self._chunk = chunk - self._served = False - - def get(self): - if self._served: - raise _StopWorker - self._served = True - return self._chunk - - -class _SettledFailureReq: - """Transfer info whose first use raises a settled transport error.""" - - room = 7 - is_dummy = False - endpoint = "10.0.0.9" - dst_port = 6009 - - @property - def agent_name(self): - raise RuntimeError("NIXL transfer encountered ERR") - - -class TestNixlWorkerSettledFailureReleasesAck(CustomTestCase): - """Bug regression: a settled transfer failure (every handle settled, decode - told via conclude_failure) left the chunk counted in _staging_outstanding, - so the abort ack for the room could never fire and the decode's deferred - KV release always ran out the full timeout instead of draining.""" - - def test_settled_failure_uncounts_chunk_and_releases_held_ack(self): - mgr = _prefill_mgr(NixlKVManager) - mgr._staging_outstanding = defaultdict(int) - mgr.enable_staging = False - mgr.exceptions = {} - mgr.decode_kv_args_table = {} - mgr.request_status[7] = KVPoll.WaitingForInput - mgr.check_status = lambda r: mgr.request_status[r] - mgr.update_status = lambda r, s: mgr.request_status.__setitem__(r, s) - mgr.record_failure = MagicMock() - mgr._await_handles = lambda handles, failure_seen=False: (True, True) - mgr.transfer_infos[7] = {"sess0": _SettledFailureReq()} - # The decode learns of the failure and its ABORT lands while the chunk - # is still counted -- the interleaving that held the ack forever. - mgr.conclude_failure = MagicMock( - side_effect=lambda **kw: mgr._handle_abort_notification( - [b"ABORT", b"7", b"10.0.0.9", b"6009"] - ) + manager.conclude_failure.assert_called_once() + self.assertEqual(manager._staging_outstanding[ROOM], 0) + self.assertNotIn(ROOM, manager._deferred_ack_poisoned_rooms) + self.assertEqual( + manager._sent, + [(ROOM, AckTarget(DECODE_IP, DECODE_PORT, ABORT_GENERATION))], ) - chunk = SimpleNamespace(room=7, staging_counted=False) + self.assertEqual(manager._deferred_ack_targets, {}) - with self.assertRaises(_StopWorker): - mgr.transfer_worker(_OneChunkQueue(chunk)) + def test_feature_off_marks_failed_without_ack(self): + manager = _manager(enabled=False) - mgr.conclude_failure.assert_called_once() - self.assertEqual(mgr._staging_outstanding[7], 0) - self.assertEqual(mgr._sent, [("10.0.0.9", 6009, 7)]) - self.assertNotIn(7, mgr._deferred_ack_targets) + claimed = manager._handle_abort_notification(self._abort_message()) + self.assertTrue(claimed) + self.assertEqual(manager.request_status[ROOM], KVPoll.Failed) + self.assertEqual(manager._deferred_ack_targets, {}) + self.assertEqual(manager._sent, []) -class TestNixlDecodeAckIngest(CustomTestCase): - def test_abort_ack_is_aggregated_per_rank(self): - # Mirrors the decode listener thread's ABORT_ACK branch. - mgr = CommonKVManager.__new__(CommonKVManager) - mgr._deferred_abort_ack_tracker = {} - mgr.register_deferred_abort_room(21) + def test_legacy_abort_without_return_address_is_tolerated(self): + manager = _manager() - for rank in (b"0", b"1", b"1"): - msg = [b"ABORT_ACK", b"21", rank] - mgr.note_abort_ack(int(msg[1].decode()), int(msg[2].decode())) + claimed = manager._handle_abort_notification( + [ABORT_TAG, str(ROOM).encode("ascii")] + ) - self.assertFalse(mgr.is_abort_release_safe(21, required_acks=3)) - self.assertTrue(mgr.is_abort_release_safe(21, required_acks=2)) + self.assertTrue(claimed) + self.assertEqual(manager.request_status[ROOM], KVPoll.Failed) + self.assertEqual(manager._deferred_ack_targets, {}) + self.assertEqual(manager._sent, []) if __name__ == "__main__": diff --git a/test/registered/unit/disaggregation/test_nixl_sender_failure_cleanup.py b/test/registered/unit/disaggregation/test_nixl_sender_failure_cleanup.py index c3078d5a3614..91d6a33b5b56 100644 --- a/test/registered/unit/disaggregation/test_nixl_sender_failure_cleanup.py +++ b/test/registered/unit/disaggregation/test_nixl_sender_failure_cleanup.py @@ -22,12 +22,20 @@ def test_failure_exception_cleans_room_state_before_raising(self): prefetched_rooms={room, 8}, prefetch_requested={(room, 0, "session-a"), (8, 0, "session-b")}, ) + # This stub mirrors CommonKVManager's field set by hand, so any new + # field touched by CommonKVSender.clear() fails here with an + # AttributeError instead of a meaningful assertion. If that keeps + # happening, spy on clear() here and test its contents where clear() + # is tested, or build a real manager instance. sender.kv_mgr = SimpleNamespace( enable_staging=True, _staging_ctx=staging_ctx, request_status={room: object()}, req_to_decode_prefix_len={room: 3}, transfer_infos={room: object()}, + _staging_outstanding={}, + _deferred_ack_targets={}, + _maybe_ack_drained_abort=lambda room: None, exceptions={room: expected_exc}, failure_records={room: "transfer failed"}, failure_lock=threading.Lock(),