From 4b8c71bb56fb07831641cfbc73374f76e62a191c Mon Sep 17 00:00:00 2001 From: Samuel Nordmann Date: Thu, 17 Sep 2026 15:16:24 +0200 Subject: [PATCH 1/2] [PD] Add PyTorch transfer and lease-lifetime review prototype Snapshot of the current endpoint-transfer prototype for review in the author's fork. Native GPU performance parity and full end-to-end framework validation remain open; this is not an upstream-ready change. Assisted-by: OpenAI Codex Signed-off-by: Samuel Nordmann --- python/sglang/srt/disaggregation/base/conn.py | 5 +- .../sglang/srt/disaggregation/common/conn.py | 398 ++- .../disaggregation/common/staging_handler.py | 221 +- .../sglang/srt/disaggregation/common/utils.py | 2 +- python/sglang/srt/disaggregation/decode.py | 221 +- .../srt/disaggregation/mooncake/conn.py | 608 +++-- python/sglang/srt/disaggregation/nixl/conn.py | 1095 ++++++-- .../srt/disaggregation/nixl/torch_transfer.py | 2169 ++++++++++++++++ python/sglang/srt/disaggregation/prefill.py | 8 + python/sglang/srt/environ.py | 10 +- python/sglang/srt/managers/io_struct.py | 4 + python/sglang/srt/managers/schedule_batch.py | 4 + python/sglang/srt/managers/scheduler.py | 18 +- .../sglang/srt/managers/tokenizer_manager.py | 4 + rust/sglang-server/src/message/io_struct.rs | 54 +- rust/sglang-server/src/message/request.rs | 30 + .../test_deferred_decode_kv_release.py | 366 ++- .../test_disaggregation_wire.py | 17 +- .../test_mooncake_generation_control.py | 192 ++ .../disaggregation/test_nixl_backend_basic.py | 846 +++++- .../test_nixl_deferred_kv_release.py | 276 +- .../test_nixl_sender_failure_cleanup.py | 12 +- .../test_nixl_torch_transfer.py | 2257 +++++++++++++++++ .../unit/managers/test_io_struct.py | 28 + 24 files changed, 8113 insertions(+), 732 deletions(-) create mode 100644 python/sglang/srt/disaggregation/nixl/torch_transfer.py create mode 100644 test/registered/unit/disaggregation/test_mooncake_generation_control.py create mode 100644 test/registered/unit/disaggregation/test_nixl_torch_transfer.py diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 358ff95c4304..e97af5a9d1ac 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -104,7 +104,10 @@ class KVPoll: class BaseKVManager(ABC): """Base class for managing transfer states""" - enable_deferred_decode_kv_release: bool = False + # Decode destination memory is always quarantined after a transfer abort + # until every possible remote writer proves quiescence. Concrete managers + # must not make this safety property optional. + enable_deferred_decode_kv_release: bool = True @abstractmethod def __init__( diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 39763024a31b..5c2439dbb52f 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -3,11 +3,13 @@ import asyncio import concurrent.futures import dataclasses +import hashlib import logging +import secrets import threading import time from collections import defaultdict -from typing import Dict, List, Optional, Set, Tuple, Union +from typing import Dict, FrozenSet, List, Optional, Set, Tuple, Union import numpy as np import numpy.typing as npt @@ -50,6 +52,57 @@ logger = logging.getLogger(__name__) +_TRANSFER_GENERATION_NBYTES = 16 + + +def validate_transfer_generation(value: object) -> str: + """Return one canonical, wire-safe 128-bit decode-room lease identity.""" + + if type(value) is not str or len(value) != 2 * _TRANSFER_GENERATION_NBYTES: + raise ValueError("transfer generation must be a 32-character string") + if value.lower() != value: + raise ValueError("transfer generation must use canonical lowercase hex") + try: + decoded = bytes.fromhex(value) + except ValueError as exc: + raise ValueError( + "transfer generation must use canonical lowercase hex" + ) from exc + if len(decoded) != _TRANSFER_GENERATION_NBYTES: + raise ValueError("transfer generation must encode exactly 128 bits") + return value + + +def next_transfer_generation(value: object) -> str: + """Derive the next shared identity for a true-retraction rebootstrap. + + The initial value is minted randomly at ingress. Chaining a domain-separated + 128-bit BLAKE2 digest lets every decode TP rank independently derive the + same fresh lease without reusing the retired generation or coordinating a + new random value through the data plane. + """ + + generation = validate_transfer_generation(value) + return hashlib.blake2b( + bytes.fromhex(generation), + digest_size=_TRANSFER_GENERATION_NBYTES, + person=b"sgl-pd-reboot-v1", + ).hexdigest() + + +@dataclasses.dataclass(frozen=True) +class _AbortAckTarget: + decode_ip: str + decode_port: int + + +@dataclasses.dataclass +class _DeferredAbortTracker: + """Exact, immutable writer set and the ACKs adopted for one room lease.""" + + expected_prefill_ranks: FrozenSet[int] + acked_prefill_ranks: Set[int] = dataclasses.field(default_factory=set) + # 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 @@ -164,9 +217,10 @@ def __init__( self.is_hybrid_mla_backend = getattr(args, "is_hybrid_mla_backend", False) self.disaggregation_mode = disaggregation_mode self.server_args = server_args - self.enable_deferred_decode_kv_release = ( - envs.SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE.get() - ) + # Compatibility/readback only. Correctness paths no longer branch on + # this value, so even an in-process mutation cannot bypass quarantine. + self.enable_deferred_decode_kv_release = True + self._abort_state_lock = threading.Lock() self._dcp_pack_buffers = None # for p/d multi node infer self.bootstrap_host = get_serving().host @@ -243,7 +297,12 @@ def __init__( 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]] = {} + self._deferred_ack_targets: Dict[Tuple[int, str], _AbortAckTarget] = {} + self._active_transfer_generations: Dict[int, str] = {} + # Without ACK-of-ACK, eviction could discard the only proof needed + # to replay a lost-return ACK. These cold-path tombstones therefore + # remain for the process lifetime. + self._abort_ack_receipts: Set[Tuple[int, str]] = set() self.req_to_decode_prefix_len: Dict[int, int] = {} self.decode_kv_args_table = {} self.pp_group = get_pp_group() @@ -265,7 +324,9 @@ def __init__( # 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]] = {} + self._deferred_abort_ack_tracker: Dict[ + Tuple[int, str], _DeferredAbortTracker + ] = {} # Heartbeat interval should be at least 2 seconds self.heartbeat_interval = max( envs.SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL.get(), 2.0 @@ -383,27 +444,145 @@ def record_failure(self, bootstrap_room: int, failure_reason: str): with self.failure_lock: self.failure_records[bootstrap_room] = failure_reason - 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, + transfer_generation: str, + expected_prefill_ranks: FrozenSet[int], + ) -> None: + """Publish the exact abort intent before any ABORT send can complete.""" - 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) + key = (int(bootstrap_room), validate_transfer_generation(transfer_generation)) + expected = frozenset(expected_prefill_ranks) + if not expected or any(type(rank) is not int or rank < 0 for rank in expected): + raise ValueError( + "expected prefill ranks must be non-empty non-negative ints" + ) + with self._abort_state_lock: + tracker = self._deferred_abort_ack_tracker.get(key) + if tracker is None: + self._deferred_abort_ack_tracker[key] = _DeferredAbortTracker(expected) + elif tracker.expected_prefill_ranks != expected: + raise RuntimeError("abort intent cannot change its expected writer set") + + def note_abort_ack( + self, bootstrap_room: int, transfer_generation: str, prefill_rank: int + ) -> bool: + """Adopt an ACK only for the exact currently quarantined room lease.""" - 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 - ) + try: + generation = validate_transfer_generation(transfer_generation) + except ValueError: + return False + if type(prefill_rank) is not int or prefill_rank < 0: + return False + key = (int(bootstrap_room), generation) + with self._abort_state_lock: + tracker = self._deferred_abort_ack_tracker.get(key) + if tracker is None or prefill_rank not in tracker.expected_prefill_ranks: + return False + tracker.acked_prefill_ranks.add(prefill_rank) + return True + + def is_abort_release_safe( + self, + bootstrap_room: int, + transfer_generation: str, + ) -> bool: + """True only after every possible writer ACKs this exact lease.""" + + key = (int(bootstrap_room), validate_transfer_generation(transfer_generation)) + with self._abort_state_lock: + tracker = self._deferred_abort_ack_tracker.get(key) + return bool( + tracker is not None + and tracker.acked_prefill_ranks == tracker.expected_prefill_ranks + ) + + def clear_deferred_abort_state( + self, bootstrap_room: int, transfer_generation: str + ) -> None: + key = (int(bootstrap_room), validate_transfer_generation(transfer_generation)) + with self._abort_state_lock: + self._deferred_abort_ack_tracker.pop(key, None) + + def register_transfer_generation( + self, bootstrap_room: int, transfer_generation: str + ) -> bool: + """Atomically bind fresh metadata, rejecting retired or active ABA.""" + + room = int(bootstrap_room) + generation = validate_transfer_generation(transfer_generation) + with self._abort_state_lock: + if (room, generation) in self._abort_ack_receipts: + return False + current = self._active_transfer_generations.get(room) + if current is not None and current != generation: + return False + self._active_transfer_generations[room] = generation + return True - def clear_deferred_abort_state(self, bootstrap_room: int) -> None: - self._deferred_abort_ack_tracker.pop(bootstrap_room, None) + def active_transfer_generation(self, bootstrap_room: int) -> Optional[str]: + with self._abort_state_lock: + return self._active_transfer_generations.get(int(bootstrap_room)) + + def _remember_abort_receipt_unlocked(self, key: Tuple[int, str]) -> None: + self._abort_ack_receipts.add(key) + + def mark_transfer_quiescent( + self, bootstrap_room: int, transfer_generation: Optional[str] = None + ) -> bool: + """Journal exact remote-writer quiescence before room teardown. + + The writer-count test and receipt publication share one critical + section with writer admission. A caller can therefore never retire a + generation between a sibling worker's increment and its DMA. + """ + + room = int(bootstrap_room) + with self._abort_state_lock: + if self._staging_outstanding.get(room, 0) > 0: + return False + generation = transfer_generation or self._active_transfer_generations.get( + room + ) + if generation is None: + return False + generation = validate_transfer_generation(generation) + key = (room, generation) + self._remember_abort_receipt_unlocked(key) + if self._active_transfer_generations.get(room) == generation: + self._active_transfer_generations.pop(room, None) + self._maybe_ack_drained_abort(room, generation) + return True + + def _count_transfer_writer(self, room: int) -> int: + """Admit one dequeued writer into the exact shared drain count.""" + + room = int(room) + with self._abort_state_lock: + count = self._staging_outstanding.get(room, 0) + 1 + self._staging_outstanding[room] = count + return count + + def _uncount_transfer_writer(self, room: int) -> int: + """Retire exactly one writer, never a concurrently active sibling.""" + + room = int(room) + with self._abort_state_lock: + count = self._staging_outstanding.get(room, 0) + if count <= 0: + raise RuntimeError(f"transfer-writer count underflow for room {room}") + remaining = count - 1 + if remaining: + self._staging_outstanding[room] = remaining + else: + self._staging_outstanding.pop(room, None) + return remaining + + def _transfer_writer_count(self, room: int) -> int: + with self._abort_state_lock: + return self._staging_outstanding.get(int(room), 0) def _prefill_unique_rank(self) -> int: """Stable per-sender id, matching what the transfer worker syncs on Success.""" @@ -413,8 +592,14 @@ def _prefill_unique_rank(self) -> int: + self.attn_cp_rank ) - def _send_abort_ack(self, decode_ip: str, decode_port: int, room: int) -> None: - """Best-effort ack that this rank's transfer for an aborted room drained.""" + def _send_abort_ack( + self, + decode_ip: str, + decode_port: int, + room: int, + transfer_generation: str, + ) -> bool: + """Send an idempotent generation-keyed ACK; report normal return only.""" try: na = NetworkAddress(decode_ip, decode_port) self._send_multipart_locked( @@ -423,28 +608,82 @@ def _send_abort_ack(self, decode_ip: str, decode_port: int, room: int) -> None: b"ABORT_ACK", str(room).encode("ascii"), str(self._prefill_unique_rank()).encode("ascii"), + transfer_generation.encode("ascii"), ], is_ipv6=na.is_ipv6, ) + return True except Exception as e: - logger.debug(f"Failed to send drained ABORT_ACK for room {room}: {e}") + logger.debug( + "Failed to send drained ABORT_ACK for room %s generation %s: %s", + room, + transfer_generation, + e, + ) + return False - 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 not None: - self._send_abort_ack(target[0], target[1], room) + def _maybe_ack_drained_abort( + self, room: int, transfer_generation: Optional[str] = None + ) -> None: + """Replay a durable ACK only after exact-generation quiescence.""" + with self._abort_state_lock: + if self._staging_outstanding.get(int(room), 0) > 0: + return + if transfer_generation is None: + keys = tuple( + key for key in self._deferred_ack_targets if key[0] == int(room) + ) + else: + generation = validate_transfer_generation(transfer_generation) + keys = ((int(room), generation),) + sends = [] + for key in keys: + target = self._deferred_ack_targets.get(key) + if target is None: + continue + if key not in self._abort_ack_receipts: + active = self._active_transfer_generations.get(int(room)) + status = self.request_status.get(int(room)) + if active != key[1] or status != KVPoll.Failed: + continue + self._remember_abort_receipt_unlocked(key) + self._active_transfer_generations.pop(int(room), None) + sends.append((key, target)) + for key, target in sends: + if self._send_abort_ack( + target.decode_ip, + target.decode_port, + key[0], + key[1], + ): + with self._abort_state_lock: + if self._deferred_ack_targets.get(key) == target: + self._deferred_ack_targets.pop(key, None) def register_deferred_ack_target( - self, room: int, decode_ip: str, decode_port: int - ) -> 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.""" - self._deferred_ack_targets[room] = (decode_ip, decode_port) + self, + room: int, + transfer_generation: str, + decode_ip: str, + decode_port: int, + ) -> bool: + """Journal a target only for an active or already-quiescent exact lease.""" + + room = int(room) + generation = validate_transfer_generation(transfer_generation) + key = (room, generation) + target = _AbortAckTarget(str(decode_ip), int(decode_port)) + with self._abort_state_lock: + if ( + self._active_transfer_generations.get(room) != generation + and key not in self._abort_ack_receipts + ): + return False + current_target = self._deferred_ack_targets.get(key) + if current_target is not None and current_target != target: + return False + self._deferred_ack_targets[key] = target + return True def get_kv_replica_factor(self) -> int: if self._kv_replica_factor is None: @@ -1341,10 +1580,8 @@ def clear(self) -> None: 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"): - # Drop a held ack target if the room concluded without draining - # (e.g. aborted before any chunk enqueued); else it leaks on prefill. - self.kv_mgr._deferred_ack_targets.pop(self.bootstrap_room, None) + # A generation-keyed ACK target is a durable ownership journal. Only a + # normally returned ACK send may remove it; sender clear is not proof. def abort(self): self.kv_mgr.record_failure( @@ -1367,8 +1604,18 @@ def __init__( mgr: CommonKVManager, bootstrap_addr: str, bootstrap_room: Optional[int] = None, + transfer_generation: Optional[str] = None, ): self.bootstrap_room = bootstrap_room + # Current request plumbing supplies one value before scheduler fanout, + # so every decode rank targeting the same prefill manager uses one + # lease identity. The fallback preserves direct legacy construction. + generation = ( + secrets.token_hex(_TRANSFER_GENERATION_NBYTES) + if transfer_generation is None + else validate_transfer_generation(transfer_generation) + ) + self._transfer_generation = generation self.bootstrap_addr = bootstrap_addr self.kv_mgr = mgr self.conclude_state: Optional[KVPoll] = None @@ -1379,6 +1626,12 @@ def __init__( self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping) + @property + def transfer_generation(self) -> str: + """Immutable generation shared by all ranks of this room lease.""" + + return self._transfer_generation + def init(self, prefill_dp_rank: int): if self.bootstrap_addr not in self.kv_mgr.prefill_info_table: self.kv_mgr.record_failure( @@ -1437,6 +1690,15 @@ def _setup_bootstrap_infos(self): target_pp_rank, ) if bootstrap_info is not None: + bootstrap_info["prefill_unique_rank"] = ( + target_tp_rank + * ( + self.prefill_info.pp_size + * self.prefill_info.attn_cp_size + ) + + target_pp_rank * self.prefill_info.attn_cp_size + + target_cp_rank + ) if self.kv_mgr.is_mla_backend: # For MLA: target_tp_rank is the selected real rank, others are dummy ranks bootstrap_info["is_dummy"] = not bool( @@ -1608,13 +1870,9 @@ def _check_waiting_timeout(self) -> Optional[KVPoll]: ) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) self.invalidate_cached_bootstrap_infos() - if ( - not self.abort_notified - and hasattr(self, "bootstrap_infos") - and self.bootstrap_infos is not None - ): + if hasattr(self, "bootstrap_infos") and self.bootstrap_infos is not None: + self.arm_abort_intent() self._send_abort_notification() - self.abort_notified = True return KVPoll.Failed def clear(self) -> None: @@ -1623,19 +1881,52 @@ def clear(self) -> None: self.kv_mgr.prefill_response_tracker.pop(self.bootstrap_room, None) def abort(self): + self.arm_abort_intent() self.kv_mgr.record_failure( self.bootstrap_room, "Aborted by AbortReq.", ) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) self.conclude_state = KVPoll.Failed + if hasattr(self, "bootstrap_infos") and self.bootstrap_infos is not None: + self._send_abort_notification() + + def arm_abort_intent(self) -> None: + """Publish quarantine ownership before the first fallible wire send.""" + + if self.abort_notified: + return + bootstrap_infos = getattr(self, "bootstrap_infos", None) + if not bootstrap_infos: + # Without rank discovery there is no authenticated writer set. Any + # deferred owner therefore remains quarantined and cannot become + # release-safe from a fabricated empty set. + logger.error( + "Cannot arm ABORT for room %s generation %s before exact " + "prefill-rank discovery; decode pages will remain quarantined", + self.bootstrap_room, + self.transfer_generation, + ) + return + expected_prefill_ranks = frozenset( + bootstrap_info["prefill_unique_rank"] for bootstrap_info in bootstrap_infos + ) + self.kv_mgr.register_deferred_abort_room( + self.bootstrap_room, + self.transfer_generation, + expected_prefill_ranks, + ) + self.abort_notified = True + + def retry_abort_notification(self) -> None: + """Replay an armed ABORT; duplicate delivery is generation-idempotent.""" + if ( - not self.abort_notified + self.abort_notified and hasattr(self, "bootstrap_infos") and self.bootstrap_infos is not None ): self._send_abort_notification() - self.abort_notified = True def _send_abort_notification(self): for bootstrap_info in self.bootstrap_infos: @@ -1649,6 +1940,7 @@ def _send_abort_notification(self): str(self.bootstrap_room).encode("ascii"), self.kv_mgr.local_ip.encode("ascii"), str(self.kv_mgr.rank_port).encode("ascii"), + self.transfer_generation.encode("ascii"), ] ) logger.debug( diff --git a/python/sglang/srt/disaggregation/common/staging_handler.py b/python/sglang/srt/disaggregation/common/staging_handler.py index a9ae31b850d5..f68f50a347a1 100644 --- a/python/sglang/srt/disaggregation/common/staging_handler.py +++ b/python/sglang/srt/disaggregation/common/staging_handler.py @@ -17,6 +17,7 @@ import torch +from sglang.srt.disaggregation.common.conn import validate_transfer_generation from sglang.srt.runtime_context import ( get_schedule, ) @@ -26,6 +27,10 @@ # Bounded wait for a watermark advance before re-enqueueing a deferred staging # chunk, so the re-enqueue retry does not busy-spin a core. STAGING_WATERMARK_WAIT_S = 0.001 +STAGING_REQ_WIRE_TAG = b"STAGING_REQ_V2" +STAGING_RSP_WIRE_TAG = b"STAGING_RSP_V2" +_MAX_STAGING_INT_DIGITS = 20 +_MAX_STAGING_SESSION_BYTES = 4096 if TYPE_CHECKING: from sglang.srt.disaggregation.decode import DecodeRequest @@ -99,8 +104,10 @@ def __init__( # before unregister runs, but release_room still needs it. self._room_to_receiver: dict = {} self._wm_subscribers: dict = {} - # room -> chunk_idx -> [(page_start, num_pages, writer_id)] fan-in - # arrivals; handler-owned so room teardown can purge them. + # room -> chunk_idx -> {exact_writer_id: (page_start, num_pages)} fan-in; + # handler-owned so room teardown can purge it. A map is required here: + # duplicate completion traffic from one rank must not impersonate the + # other writers needed to authorize scatter. self._writer_counts: dict = {} def register_wm_subscriber(self, receiver, session_id: str) -> None: @@ -307,11 +314,37 @@ def handle_chunk_arrived( ) return False room_counts = self._writer_counts.setdefault(room, {}) - arrivals = room_counts.setdefault(chunk_idx, []) - arrivals.append((page_start, num_pages, writer_id)) + arrivals = room_counts.setdefault(chunk_idx, {}) + arrival = (page_start, num_pages) + previous = arrivals.get(writer_id) + if previous is not None: + if previous != arrival: + logger.error( + "Conflicting duplicate staging completion room=%s chunk=%s " + "writer=%s prior=%s current=%s", + room, + chunk_idx, + writer_id, + previous, + arrival, + ) + return False + arrivals[writer_id] = arrival num_writers = self.num_writers_for(receiver) if len(arrivals) >= num_writers: - self.submit_chunk_scatter(room, chunk_idx, page_start, num_pages) + layouts = set(arrivals.values()) + if len(layouts) != 1: + logger.error( + "Conflicting staging completion layouts room=%s chunk=%s: %s", + room, + chunk_idx, + layouts, + ) + return False + scatter_page_start, scatter_num_pages = next(iter(layouts)) + self.submit_chunk_scatter( + room, chunk_idx, scatter_page_start, scatter_num_pages + ) del room_counts[chunk_idx] return True return False @@ -502,27 +535,99 @@ def handle_watermark_msg(staging_ctx, msg_parts) -> None: staging_ctx.watermark_cv.notify_all() -def handle_staging_rsp(msg_parts, transfer_infos: dict) -> None: - """Process a STAGING_RSP message and update transfer info with allocation.""" - stg_room = int(msg_parts[1].decode("ascii")) - stg_chunk_idx = int(msg_parts[2].decode("ascii")) - stg_offset = int(msg_parts[3].decode("ascii")) - stg_round = int(msg_parts[4].decode("ascii")) - stg_end = int(msg_parts[5].decode("ascii")) - stg_session = msg_parts[6].decode("ascii") +def _parse_staging_int(frame, field: str, *, minimum: int = 0) -> int: + if type(frame) is not bytes: + raise ValueError(f"{field} must be exact bytes") + value = frame.decode("ascii") + digits = value[1:] if value.startswith("-") else value + if ( + not digits + or len(digits) > _MAX_STAGING_INT_DIGITS + or not digits.isdecimal() + or (len(digits) > 1 and digits[0] == "0") + or value == "-0" + ): + raise ValueError(f"{field} must be canonical bounded decimal") + parsed = int(value) + if parsed < minimum: + raise ValueError(f"{field} must be at least {minimum}") + return parsed + + +def _parse_staging_session(frame) -> str: + if type(frame) is not bytes or not frame or len(frame) > _MAX_STAGING_SESSION_BYTES: + raise ValueError("staging session must be non-empty bounded bytes") + return frame.decode("ascii") + + +def _parse_staging_req(msg): + if len(msg) not in (6, 7) or msg[0] != STAGING_REQ_WIRE_TAG: + raise ValueError("unsupported STAGING_REQ wire layout") + generation = validate_transfer_generation(msg[1].decode("ascii")) + room = _parse_staging_int(msg[2], "staging room") + chunk_idx = _parse_staging_int(msg[3], "staging chunk") + chunk_num_pages = _parse_staging_int(msg[4], "staging chunk page count", minimum=1) + session_id = _parse_staging_session(msg[5]) + requester_pp_rank = ( + _parse_staging_int(msg[6], "requester PP rank") if len(msg) == 7 else None + ) + return ( + generation, + room, + chunk_idx, + chunk_num_pages, + session_id, + requester_pp_rank, + ) + + +def handle_staging_rsp(msg_parts, transfer_infos: dict) -> bool: + """Adopt an allocation response only for the exact active room lease.""" + try: + if len(msg_parts) != 8 or msg_parts[0] != STAGING_RSP_WIRE_TAG: + raise ValueError("unsupported STAGING_RSP wire layout") + generation = validate_transfer_generation(msg_parts[1].decode("ascii")) + stg_room = _parse_staging_int(msg_parts[2], "staging room") + stg_chunk_idx = _parse_staging_int(msg_parts[3], "staging chunk") + stg_offset = _parse_staging_int(msg_parts[4], "staging offset", minimum=-2) + stg_round = _parse_staging_int(msg_parts[5], "staging round") + stg_end = _parse_staging_int(msg_parts[6], "staging end", minimum=-1) + stg_session = _parse_staging_session(msg_parts[7]) + except (AttributeError, IndexError, TypeError, UnicodeError, ValueError) as error: + logger.warning("Dropping malformed or legacy STAGING_RSP: %s", error) + return False + room_infos = transfer_infos.get(stg_room, {}) tinfo = room_infos.get(stg_session) - if tinfo is not None: - if tinfo.staging is None: - tinfo.staging = StagingTransferInfo() - tinfo.staging.set_chunk(stg_chunk_idx, stg_offset, stg_round, stg_end) - else: + if tinfo is None: logger.warning( "STAGING_RSP RECV but tinfo=None room=%s chunk=%d session=%s", stg_room, stg_chunk_idx, stg_session, ) + return False + try: + active_generation = validate_transfer_generation(tinfo.transfer_generation) + except (AttributeError, TypeError, ValueError): + logger.warning( + "Dropping STAGING_RSP for room=%s with no authenticated active lease", + stg_room, + ) + return False + if generation != active_generation: + logger.warning( + "Dropping stale STAGING_RSP room=%s response_generation=%s " + "active_generation=%s", + stg_room, + generation, + active_generation, + ) + return False + if tinfo.staging is None: + tinfo.staging = StagingTransferInfo() + tinfo.staging.set_chunk(stg_chunk_idx, stg_offset, stg_round, stg_end) + return True # ====================================================================== @@ -768,7 +873,9 @@ def handle_staging_req( kv_buffer_tensors, room_receivers: dict, room_bootstrap: dict, -): + *, + parsed_request=None, +) -> bool: """Allocate staging for a chunk on-demand and send STAGING_RSP to prefill. Deduplicates: multiple prefill TP ranks requesting the same (room, chunk_idx) @@ -776,11 +883,18 @@ def handle_staging_req( """ from sglang.srt.disaggregation.common.staging_buffer import StagingAllocator - room = int(msg[1].decode("ascii")) - chunk_idx = int(msg[2].decode("ascii")) - chunk_num_pages = int(msg[3].decode("ascii")) - session_id = msg[4].decode("ascii") - requester_pp_rank = int(msg[5].decode("ascii")) if len(msg) > 5 else None + try: + ( + generation, + room, + chunk_idx, + chunk_num_pages, + session_id, + requester_pp_rank, + ) = parsed_request if parsed_request is not None else _parse_staging_req(msg) + except (AttributeError, IndexError, TypeError, UnicodeError, ValueError) as error: + logger.warning("Dropping malformed or legacy STAGING_REQ: %s", error) + return False if staging_allocator is None: logger.warning( @@ -788,7 +902,7 @@ def handle_staging_req( room, chunk_idx, ) - return + return False receiver = room_receivers.get(room) if receiver is None: @@ -798,7 +912,23 @@ def handle_staging_req( chunk_idx, session_id, ) - return + return False + try: + active_generation = validate_transfer_generation(receiver.transfer_generation) + except (AttributeError, TypeError, ValueError): + logger.warning( + "STAGING_REQ dropped: room=%s has no authenticated active lease", room + ) + return False + if generation != active_generation: + logger.warning( + "STAGING_REQ dropped: stale room=%s request_generation=%s " + "active_generation=%s", + room, + generation, + active_generation, + ) + return False infos = receiver.chunk_staging_infos if chunk_idx < len(infos) and infos[chunk_idx][0] >= 0: @@ -870,7 +1000,8 @@ def handle_staging_req( with lock: sock.send_multipart( [ - b"STAGING_RSP", + STAGING_RSP_WIRE_TAG, + generation.encode("ascii"), str(room).encode("ascii"), str(chunk_idx).encode("ascii"), str(offset).encode("ascii"), @@ -881,6 +1012,7 @@ def handle_staging_req( ) except Exception: pass + return True class StagingManagerMixin: @@ -898,8 +1030,18 @@ def _is_watermark_ready( return is_watermark_ready(self._staging_ctx, session_id, alloc_round, alloc_end) def _handle_staging_req(self, msg): - room = int(msg[1].decode("ascii")) - session_id = msg[4].decode("ascii") + try: + parsed_request = _parse_staging_req(msg) + except ( + AttributeError, + IndexError, + TypeError, + UnicodeError, + ValueError, + ) as error: + logger.warning("Dropping malformed or legacy STAGING_REQ: %s", error) + return + generation, room, _, _, session_id, _ = parsed_request handler = self._staging_handler assert handler is not None, ( "STAGING_REQ received before staging handler initialized" @@ -911,8 +1053,17 @@ def _handle_staging_req(self, msg): room, ) return + receiver = self._staging_ctx.room_receivers.get(room) + if receiver is None or receiver.transfer_generation != generation: + logger.warning( + "STAGING_REQ received for stale room lease room=%s generation=%s, " + "skipping", + room, + generation, + ) + return prefill_tp = decode_req.kv_receiver.prefill_info.attn_tp_size - handle_staging_req( + accepted = handle_staging_req( msg, self._staging_ctx.allocator, self.kv_args, @@ -921,10 +1072,10 @@ def _handle_staging_req(self, msg): getattr(self, "kv_buffer_tensors", None), self._staging_ctx.room_receivers, self._staging_ctx.room_bootstrap, + parsed_request=parsed_request, ) - receiver = self._staging_ctx.room_receivers.get(room) - if receiver is not None: + if accepted: handler.register_wm_subscriber(receiver, session_id) @@ -965,7 +1116,8 @@ def prefetch_staging_reqs( num_chunks = (total_pages + full_chunk_pages - 1) // full_chunk_pages for chunk_idx in range(num_chunks): - stg_key = (room, chunk_idx, session_id) + generation = validate_transfer_generation(tinfo.transfer_generation) + stg_key = (room, generation, chunk_idx, session_id) if stg_key in staging_requested: continue staging_requested.add(stg_key) @@ -982,7 +1134,8 @@ def prefetch_staging_reqs( sock.connect(ep) prefetch_sockets[ep] = sock request = [ - b"STAGING_REQ", + STAGING_REQ_WIRE_TAG, + generation.encode("ascii"), str(room).encode("ascii"), str(chunk_idx).encode("ascii"), str(chunk_pages).encode("ascii"), diff --git a/python/sglang/srt/disaggregation/common/utils.py b/python/sglang/srt/disaggregation/common/utils.py index 1571c3d19372..d03713650738 100644 --- a/python/sglang/srt/disaggregation/common/utils.py +++ b/python/sglang/srt/disaggregation/common/utils.py @@ -32,7 +32,7 @@ class TransferKVChunk: # Set when the staging worker first counts this chunk toward the per-room # outstanding count; stays set across re-enqueue on a watermark defer. staging_counted: bool = False - # Mori early-send: CUDA event to synchronize before RDMA (optional). + # Early-send CUDA event to synchronize before reading source memory. wait_event: Optional[object] = None diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 9d6d91966357..67d17b7cc939 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -37,7 +37,12 @@ from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.disaggregation.base import KVPoll from sglang.srt.disaggregation.base.conn import StateType -from sglang.srt.disaggregation.common.conn import CommonKVManager, CommonKVReceiver +from sglang.srt.disaggregation.common.conn import ( + CommonKVManager, + CommonKVReceiver, + next_transfer_generation, + validate_transfer_generation, +) from sglang.srt.disaggregation.decode_hicache_mixin import ( DecodeHiCachePreallocMixin, DecodeHiCacheTransferMixin, @@ -631,6 +636,35 @@ def add( if self._check_if_req_exceed_kv_capacity(req): return + backend = ( + TransferBackend.FAKE if _is_fake_transfer(req) else self.transfer_backend + ) + if not is_retracted and backend in ( + TransferBackend.NIXL, + TransferBackend.MOONCAKE, + ): + try: + generation = validate_transfer_generation(req.transfer_generation) + except ValueError as error: + message = ( + "Missing or invalid internal transfer generation on a " + f"disaggregated request: {error}" + ) + logger.error(message) + prepare_abort( + req, + message, + status_code=HTTPStatus.INTERNAL_SERVER_ERROR, + ) + self.scheduler.output_streamer.stream_output([req], req.return_logprob) + return + if is_rebootstrap: + # The completed lease has a process-lifetime tombstone and its + # delayed notifications must never match this new transfer. + # All TP ranks start from the same ingress generation, so this + # deterministic chain advances them to the same next identity. + req.transfer_generation = next_transfer_generation(generation) + if is_retracted: req.retraction_mb_id = None self.retracted_queue.append(req) @@ -699,10 +733,16 @@ def _create_receiver_and_enqueue( ) kv_receiver_class = get_kv_class(backend, KVClassType.RECEIVER) + receiver_kwargs = ( + {"transfer_generation": req.transfer_generation} + if backend in (TransferBackend.NIXL, TransferBackend.MOONCAKE) + else {} + ) kv_receiver = kv_receiver_class( mgr=self.kv_manager, bootstrap_addr=_bootstrap_addr(req), bootstrap_room=req.bootstrap_room, + **receiver_kwargs, ) decode_req = DecodeRequest( @@ -2029,6 +2069,19 @@ def _generate_fake_prefill_handoff_output_id(req: Req) -> int: return int.from_bytes(digest, byteorder="little") % req.vocab_size +@dataclass +class _DeferredKVRelease: + decode_req: DecodeRequest + bootstrap_room: int + kv_receiver: CommonKVReceiver + next_diagnostic_at: float + metadata_buffer_index: int + transfer_generation: str + diagnostic_count: int = 0 + release_phase: str = "QUARANTINED" + release_error: Optional[Tuple[str, str]] = None + + class DecodeTransferQueue(DecodeHiCacheTransferMixin): """ Store the requests that is polling kv @@ -2053,15 +2106,13 @@ def __init__( self.spec_algorithm = scheduler.spec_algorithm self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() self.staging_handler = None - self.enable_deferred_kv_release = ( - envs.SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE.get() - ) + # Non-bypassable: a remote writer must prove quiescence before these + # pages or this request slot can be reused. + self.enable_deferred_kv_release = True self.deferred_kv_release_timeout = ( envs.SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE_TIMEOUT.get() ) - # Aborted-mid-transfer requests whose KV pages/slot are held until drained - # or timed out. Entries: (decode_req, deadline, metadata_idx, required_acks). - self._deferred_releases: List[Tuple[DecodeRequest, float, int, int]] = [] + self._deferred_releases: List[_DeferredKVRelease] = [] def add(self, decode_req: DecodeRequest) -> None: self.queue.append(decode_req) @@ -2333,13 +2384,9 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req ) if self.scheduler.enable_hisparse: self.scheduler.hisparse_coordinator.request_finished(decode_req.req) - if ( - self.enable_deferred_kv_release - and decode_req.kv_receiver.kv_mgr.enable_deferred_decode_kv_release - and decode_req.kv_receiver.abort_notified - ): + if decode_req.kv_receiver.abort_notified: # Decode-initiated abort: a prefill write may still target - # these pages, so hold them until the drain ack or timeout. + # these pages, so hold them until the exact drain proof. # (A prefill-initiated failure has already stopped writing -> # immediate release below.) self._defer_release(decode_req) @@ -2412,65 +2459,143 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req def _defer_release(self, decode_req: DecodeRequest) -> None: deadline = time.monotonic() + self.deferred_kv_release_timeout - # Require an ack from every notified prefill rank (dummy-proof). Snapshot - # now -- the receiver may be cleared by resolve time. - required_acks = len(decode_req.kv_receiver.bootstrap_infos) + receiver = decode_req.kv_receiver + receiver.arm_abort_intent() self._deferred_releases.append( - (decode_req, deadline, decode_req.metadata_buffer_index, required_acks) + _DeferredKVRelease( + decode_req=decode_req, + bootstrap_room=decode_req.req.bootstrap_room, + kv_receiver=receiver, + next_diagnostic_at=deadline, + metadata_buffer_index=decode_req.metadata_buffer_index, + transfer_generation=receiver.transfer_generation, + ) ) - def _do_release(self, decode_req: DecodeRequest, idx: int) -> None: - room = decode_req.req.bootstrap_room + def _do_release(self, held: _DeferredKVRelease) -> None: + """Run one fail-stop release transaction without replaying a free. + + Each phase is journaled before its fallible or non-idempotent action. If + an action commits and its return is lost, the journal blocks automatic + retry and prefers a diagnosable leak to a possible double-free. + """ + + if held.release_phase != "QUARANTINED": + raise RuntimeError( + f"unsafe deferred release replay from phase {held.release_phase}" + ) + decode_req = held.decode_req + receiver = held.kv_receiver + room = held.bootstrap_room + idx = held.metadata_buffer_index if self.enable_staging and self.staging_handler.is_staging_room(room): + held.release_phase = "STAGING_RELEASE_ENTERED" self.staging_handler.unregister_decode_req(room) # release pre-allocated kv cache, but don't insert into the tree since it's failed + held.release_phase = "KV_RELEASE_ENTERED" release_kv_cache(decode_req.req, self.tree_cache, is_insert=False) + held.release_phase = "METADATA_RESET_ENTERED" self.metadata_buffers.bootstrap_room[idx] = 0 + held.release_phase = "METADATA_INDEX_FREE_ENTERED" self.req_to_metadata_buffer_idx_allocator.free(idx) - decode_req.kv_receiver.kv_mgr.clear_deferred_abort_state(room) - decode_req.kv_receiver.clear() + held.release_phase = "CONTROL_CLEANUP_ENTERED" + receiver.kv_mgr.clear_deferred_abort_state(room, held.transfer_generation) + receiver.clear() decode_req.kv_receiver = None + held.release_phase = "DONE" def has_pending_deferred_releases(self) -> bool: return bool(self._deferred_releases) def resolve_deferred_releases(self) -> None: - """Release held requests once every prefill rank acks the drain, or the - hold times out.""" + """Release only on exact drain proof; timeout retries and diagnoses.""" if not self._deferred_releases: return now = time.monotonic() - still_held = [] - to_release = [] - for decode_req, deadline, idx, required_acks in self._deferred_releases: - room = decode_req.req.bootstrap_room - kv_mgr = decode_req.kv_receiver.kv_mgr - drained = kv_mgr.is_abort_release_safe(room, required_acks) - if not drained and now < deadline: - still_held.append((decode_req, deadline, idx, required_acks)) - else: - to_release.append((decode_req, idx, room, drained)) - # Commit the survivors before releasing so a _do_release exception can't - # leave a released entry in the list (double-free / None receiver on retry). - self._deferred_releases = still_held - for decode_req, idx, room, drained in to_release: - if not drained: + for held in tuple(self._deferred_releases): + if held.release_phase == "DONE": + self._deferred_releases = [ + entry for entry in self._deferred_releases if entry is not held + ] + continue + if held.release_phase != "QUARANTINED": + if now >= held.next_diagnostic_at: + logger.error( + "Fail-stop KV release journal retains room %s generation " + "%s in phase %s after %s; automatic retry is unsafe", + held.bootstrap_room, + held.transfer_generation, + held.release_phase, + held.release_error, + ) + held.next_diagnostic_at = now + 300.0 + continue + decode_req = held.decode_req + receiver = held.kv_receiver + room = held.bootstrap_room + kv_mgr = receiver.kv_mgr + drained = kv_mgr.is_abort_release_safe(room, held.transfer_generation) + if drained: + try: + self._do_release(held) + except BaseException as error: + try: + error_text = str(error) + except BaseException: + error_text = "" + held.release_error = ( + f"{type(error).__module__}.{type(error).__qualname__}", + error_text[:512], + ) + held.next_diagnostic_at = now + 300.0 + logger.exception( + "Deferred KV release entered fail-stop phase %s for " + "room %s generation %s; owner remains quarantined", + held.release_phase, + room, + held.transfer_generation, + ) + if not isinstance(error, Exception): + raise + else: + self._deferred_releases = [ + entry for entry in self._deferred_releases if entry is not held + ] + continue + if now >= held.next_diagnostic_at: logger.warning( - f"Deferred KV release for room {room} timed out after " - f"{self.deferred_kv_release_timeout}s without a full drain " - f"ack from prefill; releasing anyway." + "Deferred KV quarantine for room %s generation %s has no " + "full drain proof after %.1fs; retaining pages and retrying " + "ABORT (diagnostic %d)", + room, + held.transfer_generation, + self.deferred_kv_release_timeout, + held.diagnostic_count + 1, ) - try: - self._do_release(decode_req, idx) - except Exception: - # Isolate a failed release so the rest still run; entry already dropped. - logger.exception(f"Deferred KV release failed for room {room}") + try: + receiver.retry_abort_notification() + except Exception: + logger.exception( + "ABORT retry failed for quarantined room %s generation %s", + room, + held.transfer_generation, + ) + held.diagnostic_count += 1 + backoff = min( + max(self.deferred_kv_release_timeout, 1.0) + * (2 ** min(held.diagnostic_count, 5)), + 300.0, + ) + held.next_diagnostic_at = now + backoff def release_memory_occupation(self): """Clean up in-flight transfers before releasing GPU memory.""" + if self._deferred_releases: + raise RuntimeError( + "Cannot release decode GPU memory while remote-writer " + f"quarantine holds {len(self._deferred_releases)} request(s)" + ) self.queue.clear() - # Pool is being torn down; drop held entries without per-request release. - self._deferred_releases.clear() def resume_memory_occupation(self): """Queues are already cleared on release; new transfers can be accepted.""" diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index d8914b736d98..be1e4fe26cab 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -8,7 +8,7 @@ import threading import time from collections import defaultdict -from typing import List, Optional, Set, Tuple, Union +from typing import FrozenSet, List, Optional, Set, Tuple, Union import numpy as np import numpy.typing as npt @@ -22,9 +22,12 @@ CommonKVReceiver, CommonKVSender, KVTransferError, + validate_transfer_generation, ) from sglang.srt.disaggregation.common.staging_handler import ( STAGING_WATERMARK_WAIT_S, + STAGING_REQ_WIRE_TAG, + STAGING_RSP_WIRE_TAG, DecodeStagingContext, PrefillStagingContext, StagingManagerMixin, @@ -75,6 +78,31 @@ logger = logging.getLogger(__name__) +MOONCAKE_CHUNK_READY_WIRE_TAG = b"CHUNK_READY_V2" +MOONCAKE_AUX_DATA_WIRE_TAG = b"AUX_DATA_V2" +MOONCAKE_TRANSFER_STATUS_WIRE_TAG = b"TRANSFER_STATUS_V2" +_MAX_MOONCAKE_CONTROL_UINT_DIGITS = 20 + + +def _parse_mooncake_control_uint( + frame: bytes, field: str, *, positive: bool = False +) -> int: + if type(frame) is not bytes: + raise ValueError(f"{field} must be exact bytes") + value = frame.decode("ascii") + if ( + not value + or len(value) > _MAX_MOONCAKE_CONTROL_UINT_DIGITS + or not value.isdecimal() + or (len(value) > 1 and value[0] == "0") + ): + raise ValueError(f"{field} must be canonical bounded unsigned decimal") + parsed = int(value) + if positive and parsed == 0: + raise ValueError(f"{field} must be positive") + return parsed + + FAILED_SESSION_RECOVERIES = Counter( "sglang:failed_session_recoveries_total", "Number of mooncake_session_ids un-blacklisted via probe.", @@ -95,6 +123,7 @@ class TransferInfo: is_dummy: bool decode_prefix_len: Optional[int] = None dst_device_kv_indices: Optional[npt.NDArray[np.int32]] = None + transfer_generation: str = "" # Note: always put the optional staging field at the final (it will be set through 'STAGING_RSP' pkg when needed) staging: Optional[StagingTransferInfo] = None @@ -128,6 +157,9 @@ def from_zmq(cls, msg: List[bytes]): if len(msg) > 9 and msg[9] != b"" else None ), + transfer_generation=validate_transfer_generation( + msg[10].decode("ascii") if len(msg) > 10 else None + ), ) @@ -205,8 +237,15 @@ def from_zmq(cls, msg: List[bytes]): ) +@dataclasses.dataclass(frozen=True) +class _MooncakeReceiveLease: + generation: str + session_id: str + expected_prefill_ranks: Optional[FrozenSet[int]] = None + + class MooncakeKVManager(StagingManagerMixin, CommonKVManager): - AUX_DATA_HEADER = b"AUX_DATA" + AUX_DATA_HEADER = MOONCAKE_AUX_DATA_WIRE_TAG def __init__( self, @@ -286,6 +325,7 @@ def __init__( daemon=True, ).start() elif self.disaggregation_mode == DisaggregationMode.DECODE: + self._receive_transfer_leases: dict[int, _MooncakeReceiveLease] = {} self._staging_ctx = DecodeStagingContext() if self.enable_staging else None if self.enable_staging: self._init_staging_allocator() @@ -343,6 +383,77 @@ def register_staging_room_bootstrap(self, room, bootstrap_infos, receiver): self._staging_ctx.room_bootstrap[room] = bootstrap_infos self._staging_ctx.room_receivers[room] = receiver + def _register_receive_transfer_lease( + self, + room: int, + generation: str, + session_id: str, + expected_prefill_ranks: Optional[FrozenSet[int]] = None, + ) -> None: + """Bind immutable decode control identity before accepting wire events.""" + + room = int(room) + generation = validate_transfer_generation(generation) + if type(session_id) is not str or not session_id: + raise ValueError("Mooncake receive lease requires a non-empty session") + expected = None + if expected_prefill_ranks is not None: + expected = frozenset(expected_prefill_ranks) + if not expected or any( + type(rank) is not int or rank < 0 for rank in expected + ): + raise ValueError( + "Mooncake receive lease ranks must be non-empty non-negative ints" + ) + candidate = _MooncakeReceiveLease(generation, session_id, expected) + with self._abort_state_lock: + current = self._receive_transfer_leases.get(room) + if current is None: + self._receive_transfer_leases[room] = candidate + self.prefill_response_tracker[room] = set() + return + if current.generation != generation or current.session_id != session_id: + raise RuntimeError( + "Mooncake room already has a different receive lease" + ) + if current.expected_prefill_ranks is not None and current != candidate: + raise RuntimeError("Mooncake receive lease rank identity is immutable") + self._receive_transfer_leases[room] = candidate + + @staticmethod + def _receive_control_matches_unlocked( + lease: Optional[_MooncakeReceiveLease], + generation: str, + prefill_rank: int, + *, + session_id: Optional[str] = None, + ) -> bool: + return bool( + lease is not None + and lease.expected_prefill_ranks is not None + and lease.generation == generation + and prefill_rank in lease.expected_prefill_ranks + and (session_id is None or lease.session_id == session_id) + ) + + def _clear_receive_transfer_lease(self, room: int, generation: str) -> bool: + """Clear only this receiver's identity and room state, atomically.""" + + room = int(room) + generation = validate_transfer_generation(generation) + with self._abort_state_lock: + current = self._receive_transfer_leases.get(room) + # Absence is not proof of ownership. A duplicate delayed clear for + # lease A can otherwise erase room state that lease B published + # immediately before B atomically binds its receive identity. + if current is None or current.generation != generation: + return False + self._receive_transfer_leases.pop(room, None) + self.request_status.pop(room, None) + self.required_prefill_response_num_table.pop(room, None) + self.prefill_response_tracker.pop(room, None) + return True + def set_kv_buffer_tensors( self, k_buffers: list, @@ -402,7 +513,8 @@ def _send_chunk_ready(self, req, chunk_idx, kv_chunk, prefill_unique_rank): self._send_multipart_locked( na.to_tcp(), [ - b"CHUNK_READY", + MOONCAKE_CHUNK_READY_WIRE_TAG, + req.transfer_generation.encode("ascii"), str(req.room).encode("ascii"), str(chunk_idx).encode("ascii"), str(kv_chunk.index_slice.start).encode("ascii"), @@ -847,9 +959,9 @@ def _validate_envelope_kv_layout( def _await_transfer_futures(self, futures) -> int: """Await a chunk's per-layer RDMA writes; return the first non-zero status. - cancel() is a no-op for a running future, so with deferred release on we - still drain the running ones before returning (no write may outlive this - call, which the drain-ack relies on). Off: original early-return.""" + cancel() is a no-op for a running future, so every running future is + drained before returning. No write may outlive the call that publishes + the drain proof, regardless of the legacy compatibility field.""" ret = 0 for future in concurrent.futures.as_completed(futures): try: @@ -860,8 +972,6 @@ def _await_transfer_futures(self, futures) -> int: ret = status for f in futures: f.cancel() - if not self.enable_deferred_decode_kv_release: - return ret return ret def send_kvcache( @@ -1160,12 +1270,13 @@ def send_aux( req: TransferInfo, prefill_aux_index: int, dst_aux_ptrs: list[int], + prefill_rank: int, ): # TODO(shangming): Fix me when nvlink_transport of Mooncake is bug-free if ( self.enable_custom_mem_pool and self.custom_mem_pool_type == "NVLINK" ) or envs.SGLANG_MOONCAKE_SEND_AUX_TCP.get(): - return self.send_aux_tcp(req, prefill_aux_index, dst_aux_ptrs) + return self.send_aux_tcp(req, prefill_aux_index, dst_aux_ptrs, prefill_rank) transfer_blocks = [] prefill_aux_ptrs = self.kv_args.aux_data_ptrs @@ -1184,6 +1295,7 @@ def send_aux_tcp( req: TransferInfo, prefill_aux_index: int, dst_aux_ptrs: list[int], + prefill_rank: int, ): prefill_aux_ptrs = self.kv_args.aux_data_ptrs prefill_aux_item_lens = self.kv_args.aux_item_lens @@ -1197,6 +1309,8 @@ def send_aux_tcp( remote=req.endpoint, dst_port=req.dst_port, room=req.room, + transfer_generation=req.transfer_generation, + prefill_rank=prefill_rank, buffer_index=i, aux_index=req.dst_aux_index, data=data, @@ -1209,6 +1323,8 @@ def send_aux_data_to_endpoint( remote: str, dst_port: int, room: int, + transfer_generation: str, + prefill_rank: int, buffer_index: int, aux_index: int, data: bytes, @@ -1218,7 +1334,9 @@ def send_aux_data_to_endpoint( na.to_tcp(), [ MooncakeKVManager.AUX_DATA_HEADER, + validate_transfer_generation(transfer_generation).encode("ascii"), str(room).encode("ascii"), + str(prefill_rank).encode("ascii"), str(buffer_index).encode("ascii"), str(aux_index).encode("ascii"), struct.pack(">I", len(data)), @@ -1227,25 +1345,62 @@ def send_aux_data_to_endpoint( is_ipv6=na.is_ipv6, ) - def _handle_aux_data(self, msg: List[bytes]): - """Handle AUX_DATA messages received by the decode thread.""" - room = int(msg[1].decode("ascii")) - buffer_index = int(msg[2].decode("ascii")) - aux_index = int(msg[3].decode("ascii")) - data_length = struct.unpack(">I", msg[4])[0] - data = msg[5] - - if len(data) != data_length: - logger.error(f"AUX_DATA length mismatch for bootstrap_room {room}") - return - - AuxDataCodec.deserialize_data_to_buffer( - self.kv_args, buffer_index, aux_index, data - ) + def _handle_aux_data(self, msg: List[bytes]) -> bool: + """Handle one exact-lease AUX_DATA message without room-reuse ABA.""" + try: + if len(msg) != 8 or msg[0] != MOONCAKE_AUX_DATA_WIRE_TAG: + raise ValueError("unsupported AUX_DATA wire layout") + generation = validate_transfer_generation(msg[1].decode("ascii")) + room = _parse_mooncake_control_uint(msg[2], "room") + prefill_rank = _parse_mooncake_control_uint(msg[3], "prefill rank") + buffer_index = _parse_mooncake_control_uint(msg[4], "buffer index") + aux_index = _parse_mooncake_control_uint(msg[5], "aux index") + if type(msg[6]) is not bytes or len(msg[6]) != 4: + raise ValueError("AUX_DATA length must be one uint32 frame") + data_length = struct.unpack(">I", msg[6])[0] + data = msg[7] + if type(data) is not bytes or len(data) != data_length: + raise ValueError("AUX_DATA payload length mismatch") + except ( + AttributeError, + IndexError, + TypeError, + UnicodeError, + ValueError, + ) as error: + logger.warning("Dropping malformed or legacy AUX_DATA: %s", error) + return False + + # Identity validation and the destination write are one critical + # section with receiver clear/rebind. An old lease can therefore never + # pass validation and then write through a newly reused room. + with self._abort_state_lock: + lease = self._receive_transfer_leases.get(room) + if not self._receive_control_matches_unlocked( + lease, generation, prefill_rank + ): + logger.warning( + "Dropping stale or unexpected-rank AUX_DATA room=%s " + "generation=%s rank=%s", + room, + generation, + prefill_rank, + ) + return False + try: + AuxDataCodec.deserialize_data_to_buffer( + self.kv_args, buffer_index, aux_index, data + ) + except Exception as error: + logger.error("Failed to apply authenticated AUX_DATA: %s", error) + return False logger.debug( - f"Received AUX_DATA for bootstrap_room {room} with length:{len(data)}" + "Received generation-bound AUX_DATA for bootstrap_room %s with length:%s", + room, + len(data), ) + return True def _get_dsa_cache_transfer_skip_flags( self, info: Optional[KVArgsRegisterInfo] @@ -1684,12 +1839,20 @@ def _send_mamba_state_slice( return self._transfer_data(req.mooncake_session_id, transfer_blocks) def sync_status_to_decode_endpoint( - self, remote: str, dst_port: int, room: int, status: int, prefill_rank: int + self, + remote: str, + dst_port: int, + room: int, + status: int, + prefill_rank: int, + transfer_generation: str, ): na = NetworkAddress(remote, dst_port) self._send_multipart_locked( na.to_tcp(), [ + MOONCAKE_TRANSFER_STATUS_WIRE_TAG, + validate_transfer_generation(transfer_generation).encode("ascii"), str(room).encode("ascii"), str(status).encode("ascii"), str(prefill_rank).encode("ascii"), @@ -1727,7 +1890,7 @@ def transfer_worker( # the predicate the abort ack relies on. The flag survives # re-enqueue on defer. if not kv_chunk.staging_counted: - self._staging_outstanding[kv_chunk.room] += 1 + self._count_transfer_writer(kv_chunk.room) kv_chunk.staging_counted = True if ( @@ -1743,10 +1906,13 @@ def transfer_worker( MooncakeRequestStage.MOONCAKE_WORKER_SEND.level, thread_finish_flag=True, ) - self._staging_outstanding.pop(kv_chunk.room, None) - if self.enable_deferred_decode_kv_release: - # Skipped => nothing written for this aborted room; ack. - self._maybe_ack_drained_abort(kv_chunk.room) + remaining_writers = self._uncount_transfer_writer(kv_chunk.room) + kv_chunk.staging_counted = False + if remaining_writers == 0 and self.mark_transfer_quiescent( + kv_chunk.room + ): + self.transfer_infos.pop(kv_chunk.room, None) + self.req_to_decode_prefix_len.pop(kv_chunk.room, None) continue if ( @@ -1788,6 +1954,7 @@ def transfer_worker( req.room, KVPoll.Failed, prefill_unique_rank, + req.transfer_generation, ) break @@ -1942,6 +2109,7 @@ def transfer_worker( req.room, KVPoll.Failed, prefill_unique_rank, + req.transfer_generation, ) break @@ -1973,6 +2141,7 @@ def transfer_worker( req.room, KVPoll.Failed, prefill_unique_rank, + req.transfer_generation, ) break @@ -1981,23 +2150,35 @@ def transfer_worker( req, kv_chunk.prefill_aux_index, target_rank_registration_info.dst_aux_ptrs, + prefill_unique_rank, ) polls.append(True if ret == 0 else False) dst_ranks_infos.append( - (req.endpoint, req.dst_port, req.room) + ( + req.endpoint, + req.dst_port, + req.room, + req.transfer_generation, + ) ) # Only sync status when all the dst ranks have received the kvcache if len(polls) == req.required_dst_info_num: status = KVPoll.Success if all(polls) else KVPoll.Failed self.update_status(req.room, status) - for endpoint, dst_port, room in dst_ranks_infos: + for ( + endpoint, + dst_port, + room, + generation, + ) in dst_ranks_infos: self.sync_status_to_decode_endpoint( endpoint, dst_port, room, status, prefill_unique_rank, + generation, ) else: # Dummy request means the decode instance is not used, so its status can be marked as success directly @@ -2022,25 +2203,22 @@ def transfer_worker( if staging_deferred: continue - self._staging_outstanding[kv_chunk.room] -= 1 - if self.enable_deferred_decode_kv_release: - # In-flight write finished; if aborted and nothing outstanding, - # the pages are idle -> release the held ack. - self._maybe_ack_drained_abort(kv_chunk.room) + remaining_writers = self._uncount_transfer_writer(kv_chunk.room) + kv_chunk.staging_counted = False # Tear down only when no chunk is still outstanding and the room # has concluded: already cleared, Success, or a Failed *last* # chunk. A non-last Failed chunk keeps the room (more chunks may # follow), not on the last chunk alone since an earlier deferred # chunk may still need to transfer. - if self._staging_outstanding.get(kv_chunk.room, 0) <= 0 and ( - kv_chunk.room not in self.request_status - or self.check_status(kv_chunk.room) == KVPoll.Success - or ( - kv_chunk.is_last_chunk - and self.check_status(kv_chunk.room) == KVPoll.Failed + if ( + remaining_writers == 0 + and ( + kv_chunk.room not in self.request_status + or self.check_status(kv_chunk.room) == KVPoll.Success + or self.check_status(kv_chunk.room) == KVPoll.Failed ) + and self.mark_transfer_quiescent(kv_chunk.room) ): - self._staging_outstanding.pop(kv_chunk.room, None) if kv_chunk.room in self.transfer_infos: self.transfer_infos.pop(kv_chunk.room) self.req_to_decode_prefix_len.pop(kv_chunk.room, None) @@ -2070,76 +2248,52 @@ def bootstrap_thread(): handle_watermark_msg(self._staging_ctx, waiting_req_bytes) continue # Staging: decode replies with allocated staging offset - if room == "STAGING_RSP": + if waiting_req_bytes[0] == STAGING_RSP_WIRE_TAG: 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")) + try: + if len(waiting_req_bytes) != 5: + raise ValueError("ABORT must contain exactly five frames") + 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")) + transfer_generation = validate_transfer_generation( + waiting_req_bytes[4].decode("ascii") + ) + except Exception as error: + logger.warning("Ignoring malformed ABORT: %s", error) + continue + generation_active = ( + self.active_transfer_generation(room_to_be_aborted) + == transfer_generation + ) room_active = ( - room_to_be_aborted in self.request_status + generation_active + and 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 nothing is in - # flight, decode falls back to the release timeout. - if room_active: - 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" - ) - elif self._staging_outstanding.get(room_to_be_aborted, 0) == 0: - # Concluded/unknown AND quiescent: ack now. A cleared - # room is not automatically quiescent -- clear() can - # drop a room whose chunk is still transferring. - 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, + if self.register_deferred_ack_target( + room_to_be_aborted, + transfer_generation, + decode_ip, + decode_port, + ): + self._maybe_ack_drained_abort( + room_to_be_aborted, transfer_generation ) + elif not generation_active: 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}" + "Ignoring stale ABORT for room %s generation %s", + room_to_be_aborted, + transfer_generation, ) continue mooncake_session_id = waiting_req_bytes[3].decode("ascii") @@ -2168,14 +2322,37 @@ def bootstrap_thread(): ) continue else: - required_dst_info_num = int(waiting_req_bytes[7].decode("ascii")) - room = int(room) + try: + required_dst_info_num = int( + waiting_req_bytes[7].decode("ascii") + ) + room = int(room) + transfer_info = TransferInfo.from_zmq(waiting_req_bytes) + except (IndexError, UnicodeError, ValueError) as error: + # A recyclable room number is not an identity. Missing + # or malformed generation metadata cannot fail whichever + # newer lease happens to own that number, and must not + # terminate the long-lived listener thread. + logger.warning( + "Dropping unauthenticated Mooncake room metadata: %s", + error, + ) + continue + if not self.register_transfer_generation( + room, transfer_info.transfer_generation + ): + logger.warning( + "Dropping late Mooncake metadata for room %s " + "generation %s; active generation is %s", + room, + transfer_info.transfer_generation, + self.active_transfer_generation(room), + ) + continue if room not in self.transfer_infos: self.transfer_infos[room] = {} - self.transfer_infos[room][mooncake_session_id] = ( - TransferInfo.from_zmq(waiting_req_bytes) - ) + self.transfer_infos[room][mooncake_session_id] = transfer_info # NOTE: after bootstrapping we can mark the req as waiting for input if len(self.transfer_infos[room]) == required_dst_info_num: self.resolve_kv_replica_factor(self.transfer_infos[room]) @@ -2200,72 +2377,152 @@ def decode_thread(): continue # Staging: prefill notifies a chunk written to staging buffer - if msg[0] == b"CHUNK_READY": - room = int(msg[1].decode("ascii")) - chunk_idx = int(msg[2].decode("ascii")) - page_start = int(msg[3].decode("ascii")) - num_pages = int(msg[4].decode("ascii")) - session_id = msg[5].decode("ascii") - handler = self._staging_handler - assert handler is not None, ( - "CHUNK_READY received before staging handler initialized" - ) - handler.handle_chunk_arrived( - room, - chunk_idx, - page_start, - num_pages, - session_id, - ) + if msg[0] == MOONCAKE_CHUNK_READY_WIRE_TAG: + self._handle_chunk_ready(msg) continue # Staging: prefill pre-requests staging allocation before forward - if msg[0] == b"STAGING_REQ": + if msg[0] == STAGING_REQ_WIRE_TAG: 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")) + try: + if len(msg) != 4: + raise ValueError( + "ABORT_ACK must contain exactly four frames" + ) + ack_aborted_room = int(msg[1].decode("ascii")) + prefill_rank = int(msg[2].decode("ascii")) + generation = validate_transfer_generation( + msg[3].decode("ascii") + ) + except (UnicodeDecodeError, ValueError): + logger.warning( + "Ignoring generation-less or malformed ABORT_ACK" ) + else: + logger.debug("Received ABORT_ACK for room %s", ack_aborted_room) + self.note_abort_ack(ack_aborted_room, generation, prefill_rank) continue - bootstrap_room, status, prefill_rank = msg - status = int(status.decode("ascii")) - bootstrap_room = int(bootstrap_room.decode("ascii")) - prefill_rank = int(prefill_rank.decode("ascii")) - - if status == KVPoll.Success: - if bootstrap_room in self.request_status: - self.prefill_response_tracker[bootstrap_room].add(prefill_rank) - expected_response_num = ( - self.required_prefill_response_num_table[bootstrap_room] - ) - arrived_response_num = len( - self.prefill_response_tracker[bootstrap_room] - ) - if arrived_response_num == expected_response_num: - if self.enable_staging: - handler = self._staging_handler - if handler.is_staging_room(bootstrap_room): - handler.submit_last_scatter_async(bootstrap_room) - self.update_status(bootstrap_room, KVPoll.Success) - elif status == KVPoll.Failed: - self.record_failure( - bootstrap_room, - "Failed to get kvcache from prefill instance, it might be dead", - ) - self.update_status(bootstrap_room, status) + if msg[0] == MOONCAKE_TRANSFER_STATUS_WIRE_TAG: + self._handle_transfer_status(msg) + continue + logger.warning( + "Dropping unknown, legacy, or malformed Mooncake control frame: %r", + msg[0][:32] if msg else b"", + ) threading.Thread(target=decode_thread).start() self._start_heartbeat_checker_thread() + def _handle_chunk_ready(self, msg: List[bytes]) -> bool: + try: + if len(msg) != 8 or msg[0] != MOONCAKE_CHUNK_READY_WIRE_TAG: + raise ValueError("unsupported CHUNK_READY wire layout") + generation = validate_transfer_generation(msg[1].decode("ascii")) + room = _parse_mooncake_control_uint(msg[2], "room") + chunk_idx = _parse_mooncake_control_uint(msg[3], "chunk index") + page_start = _parse_mooncake_control_uint(msg[4], "page start") + num_pages = _parse_mooncake_control_uint( + msg[5], "page count", positive=True + ) + session_id = msg[6].decode("ascii") + if not session_id: + raise ValueError("CHUNK_READY session must be non-empty") + prefill_rank = _parse_mooncake_control_uint(msg[7], "prefill rank") + except ( + AttributeError, + IndexError, + TypeError, + UnicodeError, + ValueError, + ) as error: + logger.warning("Dropping malformed or legacy CHUNK_READY: %s", error) + return False + + with self._abort_state_lock: + lease = self._receive_transfer_leases.get(room) + if not self._receive_control_matches_unlocked( + lease, + generation, + prefill_rank, + session_id=session_id, + ): + logger.warning( + "Dropping stale or unexpected-rank CHUNK_READY room=%s " + "generation=%s rank=%s", + room, + generation, + prefill_rank, + ) + return False + handler = self._staging_handler + if handler is None: + logger.warning("CHUNK_READY received before staging initialization") + return False + handler.handle_chunk_arrived( + room, + chunk_idx, + page_start, + num_pages, + prefill_rank, + ) + return True + + def _handle_transfer_status(self, msg: List[bytes]) -> bool: + try: + if len(msg) != 5 or msg[0] != MOONCAKE_TRANSFER_STATUS_WIRE_TAG: + raise ValueError("unsupported TRANSFER_STATUS wire layout") + generation = validate_transfer_generation(msg[1].decode("ascii")) + room = _parse_mooncake_control_uint(msg[2], "room") + status = _parse_mooncake_control_uint(msg[3], "status") + prefill_rank = _parse_mooncake_control_uint(msg[4], "prefill rank") + if status not in (KVPoll.Success, KVPoll.Failed): + raise ValueError("TRANSFER_STATUS must be Success or Failed") + except ( + AttributeError, + IndexError, + TypeError, + UnicodeError, + ValueError, + ) as error: + logger.warning("Dropping malformed or legacy TRANSFER_STATUS: %s", error) + return False + + with self._abort_state_lock: + lease = self._receive_transfer_leases.get(room) + if not self._receive_control_matches_unlocked( + lease, generation, prefill_rank + ): + logger.warning( + "Dropping stale or unexpected-rank TRANSFER_STATUS room=%s " + "generation=%s rank=%s", + room, + generation, + prefill_rank, + ) + return False + if status == KVPoll.Failed: + self.record_failure( + room, + "Failed to get kvcache from prefill instance, it might be dead", + ) + self.update_status(room, KVPoll.Failed) + return True + + self.prefill_response_tracker[room].add(prefill_rank) + if self.prefill_response_tracker[room] != lease.expected_prefill_ranks: + return True + if self.enable_staging: + handler = self._staging_handler + if handler is not None and handler.is_staging_room(room): + handler.submit_last_scatter_async(room) + self.update_status(room, KVPoll.Success) + return True + def add_transfer_request( self, bootstrap_room: int, @@ -2499,10 +2756,46 @@ def __init__( mgr: MooncakeKVManager, bootstrap_addr: str, bootstrap_room: Optional[int] = None, + transfer_generation: Optional[str] = None, ): self.session_id = mgr.get_session_id() self.init_time = None - super().__init__(mgr, bootstrap_addr, bootstrap_room) + super().__init__( + mgr, + bootstrap_addr, + bootstrap_room, + transfer_generation=transfer_generation, + ) + # Publish the generation before any delayed control can be attributed + # to the freshly created room. Rank identity is bound after bootstrap. + self.kv_mgr._register_receive_transfer_lease( + self.bootstrap_room, + self.transfer_generation, + self.session_id, + ) + + def init(self, prefill_dp_rank: int): + super().init(prefill_dp_rank) + if self.conclude_state == KVPoll.Failed: + return + expected_prefill_ranks = frozenset( + bootstrap_info["prefill_unique_rank"] + for bootstrap_info in self.bootstrap_infos + if not bootstrap_info["is_dummy"] + ) + self.kv_mgr._register_receive_transfer_lease( + self.bootstrap_room, + self.transfer_generation, + self.session_id, + expected_prefill_ranks, + ) + + def clear(self) -> None: + # The exact-generation compare and all room-keyed state removal share + # one lock; a delayed clear from lease A cannot erase lease B. + self.kv_mgr._clear_receive_transfer_lease( + self.bootstrap_room, self.transfer_generation + ) def _register_kv_args(self) -> bool: for bootstrap_info in self.bootstrap_infos: @@ -2644,6 +2937,7 @@ def send_metadata( if not is_dummy and device_kv_indices is not None else b"" ), + self.transfer_generation.encode("ascii"), ] ) except zmq.ZMQError: diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index a7fc753ad01b..463d1eb9badf 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -8,7 +8,7 @@ import time import uuid from collections import defaultdict -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple +from typing import TYPE_CHECKING, Any, Dict, FrozenSet, List, Optional, Set, Tuple import numpy as np import numpy.typing as npt @@ -24,9 +24,12 @@ CommonKVReceiver, CommonKVSender, KVTransferError, + validate_transfer_generation, ) from sglang.srt.disaggregation.common.staging_handler import ( STAGING_WATERMARK_WAIT_S, + STAGING_REQ_WIRE_TAG, + STAGING_RSP_WIRE_TAG, StagingManagerMixin, handle_staging_rsp, handle_watermark_msg, @@ -71,6 +74,126 @@ GUARD = "NixlMsgGuard".encode("ascii") KV_MEM_KINDS = {"VRAM", "DRAM"} +_COMPLETION_NOTIFICATION_VERSION = "nixlv2" +_MAX_NOTIFICATION_UINT_DIGITS = 20 + + +def _parse_notification_uint(value: str, field: str, *, positive: bool = False) -> int: + """Parse one canonical bounded unsigned integer from a completion tag.""" + + if ( + not value + or len(value) > _MAX_NOTIFICATION_UINT_DIGITS + or not value.isascii() + or not value.isdecimal() + or (len(value) > 1 and value[0] == "0") + ): + raise ValueError(f"{field} must be canonical unsigned decimal") + parsed = int(value) + if positive and parsed == 0: + raise ValueError(f"{field} must be positive") + return parsed + + +def _parse_completion_notification( + message: bytes, +) -> Tuple[str, int, str, Tuple[Any, ...]]: + """Parse the strict, generation-bound NIXL completion wire contract. + + Every accepted tag starts with ``nixlv2_{generation}_{room}_{kind}``. + Legacy/generation-less tags are intentionally incompatible: accepting one + would let a queued completion from a retired room lease complete its reuse. + """ + + if type(message) is not bytes: + raise ValueError("completion notification must be exact bytes") + components = message.decode("ascii").split("_") + if len(components) < 4 or components[0] != _COMPLETION_NOTIFICATION_VERSION: + raise ValueError("unsupported completion notification version") + generation = validate_transfer_generation(components[1]) + room = _parse_notification_uint(components[2], "room") + kind = components[3] + + if kind == "kv": + if len(components) not in (7, 10): + raise ValueError("invalid kv completion field count") + chunk_id = _parse_notification_uint(components[4], "chunk_id") + if components[5] not in ("0", "1"): + raise ValueError("is_last must be exactly 0 or 1") + is_last = components[5] == "1" + pp_rank = _parse_notification_uint(components[6], "pp_rank") + if len(components) == 7: + return generation, room, kind, (chunk_id, is_last, pp_rank) + if components[7] != "part": + raise ValueError("invalid kv-part discriminator") + part_idx = _parse_notification_uint(components[8], "part_idx") + num_parts = _parse_notification_uint(components[9], "num_parts", positive=True) + if part_idx >= num_parts: + raise ValueError("kv-part index is outside its declared part count") + return ( + generation, + room, + "kvpart", + ( + chunk_id, + is_last, + pp_rank, + part_idx, + num_parts, + ), + ) + + if kind == "state": + if len(components) not in (5, 6): + raise ValueError("invalid state completion field count") + pp_rank = _parse_notification_uint(components[4], "pp_rank") + component_id = ( + _parse_notification_uint(components[5], "state_component") + if len(components) == 6 + else None + ) + return generation, room, kind, (pp_rank, component_id) + + if kind == "aux": + if len(components) == 5: + pp_rank = _parse_notification_uint(components[4], "pp_rank") + return generation, room, kind, (pp_rank, None) + if len(components) != 7 or components[4] != "nokv": + raise ValueError("invalid aux completion layout") + pp_rank = _parse_notification_uint(components[5], "pp_rank") + expected = _parse_notification_uint(components[6], "expected_kv_chunks") + return generation, room, kind, (pp_rank, expected) + + if kind == "stg": + if len(components) < 11: + raise ValueError("invalid staging completion field count") + chunk_id = _parse_notification_uint(components[4], "chunk_id") + if components[5] not in ("0", "1"): + raise ValueError("is_last must be exactly 0 or 1") + is_last = components[5] == "1" + pp_rank = _parse_notification_uint(components[6], "pp_rank") + chunk_idx = _parse_notification_uint(components[7], "staging_chunk_idx") + page_start = _parse_notification_uint(components[8], "page_start") + num_pages = _parse_notification_uint(components[9], "num_pages", positive=True) + agent_name = "_".join(components[10:]) + if not agent_name: + raise ValueError("staging agent name must be non-empty") + return ( + generation, + room, + kind, + ( + chunk_id, + is_last, + pp_rank, + chunk_idx, + page_start, + num_pages, + agent_name, + ), + ) + + raise ValueError(f"unsupported completion notification kind: {kind!r}") def _normalize_kv_mem_kinds(kinds: Optional[List[str]], expected_len: int) -> List[str]: @@ -167,10 +290,23 @@ class TransferInfo: dst_state_indices: List[List[int]] decode_prefix_len: Optional[int] = None # for decode radix cache is_dummy_rank: Optional[bool] = None + transfer_generation: str = "" # NOTE: optional staging field; populated via STAGING_RSP. Keep at the # end so positional construction in from_zmq() continues to work. staging: Optional[StagingTransferInfo] = None + def __post_init__(self): + generation = validate_transfer_generation(self.transfer_generation) + # Cache the immutable per-lease prefix once on metadata ingestion. The + # per-chunk hot path only appends its small kind-specific suffix. + self._completion_notification_prefix = ( + f"{_COMPLETION_NOTIFICATION_VERSION}_{generation}_{self.room}" + ) + + @property + def completion_notification_prefix(self) -> str: + return self._completion_notification_prefix + def is_dummy(self): # A transfer is "dummy" only for CP non-authoritative ranks. # When dst_kv_indices is empty due to a decode-side radix cache @@ -205,6 +341,9 @@ def from_zmq(cls, msg: List[bytes]): if len(msg) > 9 and msg[9] != b"" else None ), + transfer_generation=validate_transfer_generation( + msg[10].decode("ascii") if len(msg) > 10 else None + ), ) @@ -360,6 +499,13 @@ def repeat_indices_over_layers( class TransferStatus: """Used by KV Receiver to know when a transfer is done.""" + # Exact receiver lease identity. A room number alone is recyclable and + # therefore cannot authenticate an attached completion notification. + transfer_generation: str + # Bound after bootstrap discovery and before metadata submission. Every + # completion family must name one exact non-dummy remote writer. + expected_source_ranks: Optional[FrozenSet[int]] = None + # KV chunks received per pp_rank: {pp_rank: set of chunk_ids} received_kvs_per_pp: Dict[int, Set[int]] = dataclasses.field( default_factory=lambda: defaultdict(set) @@ -379,16 +525,51 @@ class TransferStatus: received_kv_parts_per_pp: Optional[Dict[Tuple[int, int], Set[int]]] = None expected_kv_parts_per_pp: Optional[Dict[Tuple[int, int], int]] = None + def __post_init__(self): + self.transfer_generation = validate_transfer_generation( + self.transfer_generation + ) + if self.expected_source_ranks is not None: + self.bind_expected_source_ranks(self.expected_source_ranks) + + def bind_expected_source_ranks(self, ranks) -> None: + expected = frozenset(ranks) + if not expected or any(type(rank) is not int or rank < 0 for rank in expected): + raise ValueError( + "expected completion ranks must be non-empty non-negative ints" + ) + if ( + self.expected_source_ranks is not None + and self.expected_source_ranks != expected + ): + raise RuntimeError("completion source-rank identity set is immutable") + self.expected_source_ranks = expected + self.num_pp_ranks_expected = len(expected) + + def accepts_source_rank(self, rank: object) -> bool: + return ( + type(rank) is int + and self.expected_source_ranks is not None + and rank in self.expected_source_ranks + ) + def is_done(self): if self.num_pp_ranks_expected is None or not self.received_aux: return False # If state data is expected, check all PP ranks have sent it - if ( - self.expects_state - and len(self.received_state_per_pp) < self.num_pp_ranks_expected - ): - return False + if self.expects_state: + if self.expected_source_ranks is not None: + if not self.expected_source_ranks.issubset(self.received_state_per_pp): + return False + elif len(self.received_state_per_pp) < self.num_pp_ranks_expected: + # Compatibility for direct status construction outside the + # receiver path. Production status is exact-set bound in init(). + return False # All PP ranks must have reported their expected count + if self.expected_source_ranks is not None and set( + self.expected_kvs_per_pp + ) != set(self.expected_source_ranks): + return False if len(self.expected_kvs_per_pp) < self.num_pp_ranks_expected: return False # Each PP rank must have received all expected chunks @@ -407,9 +588,9 @@ def __init__( is_mla_backend: Optional[bool] = False, ): super().__init__(args, disaggregation_mode, server_args, is_mla_backend) - self.transfer_source_rank = ( - self.kv_args.pp_rank * get_parallel().tp_size + self.kv_args.engine_rank - ) + # One collision-free TP×PP×CP identity, shared with ABORT_ACK and the + # receiver's authenticated bootstrap-derived source-rank set. + self.transfer_source_rank = self._prefill_unique_rank() self.kv_args.kv_data_mem_kinds = _normalize_kv_mem_kinds( getattr(self.kv_args, "kv_data_mem_kinds", None), len(self.kv_args.kv_data_ptrs), @@ -420,15 +601,10 @@ def __init__( and self.kv_args.kv_data_mem_kinds else None ) - try: - from nixl._api import nixl_agent, nixl_agent_config, nixl_thread_sync_t - except ImportError as e: - raise ImportError( - "Please install NIXL by following the instructions at " - "https://github.com/ai-dynamo/nixl/blob/main/README.md " - "to run SGLang with NixlTransferEngine." - ) from e - + self._use_torch_transfer = ( + envs.SGLANG_DISAGGREGATION_NIXL_USE_TORCH_TRANSFER.get() + ) + self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() backend = envs.SGLANG_DISAGGREGATION_NIXL_BACKEND.get() num_threads = 8 if disaggregation_mode == DisaggregationMode.PREFILL else 0 backend_params = json.loads( @@ -442,14 +618,6 @@ def __init__( "SGLANG_DISAGGREGATION_NIXL_BACKEND_PARAMS must be a JSON object " "with string keys and string values" ) - # self.transfer_worker and self._start_bootstrap_thread runs concurrently - # so we cannot use sync_mode=None which is thread-unsafe. - agent_config = nixl_agent_config( - backends=[], - num_threads=num_threads, - sync_mode=nixl_thread_sync_t.NIXL_THREAD_SYNC_STRICT, - ) - self.agent = nixl_agent(str(uuid.uuid4()), agent_config) if num_threads > 0: # TODO: Remove this once NIXL passes thread parameters from # nixl_agent_config to explicitly-created backends. @@ -459,6 +627,55 @@ def __init__( backend_params.setdefault("thread_count", str(num_threads)) elif backend == "UCCL": backend_params.setdefault("num_cpus", str(num_threads)) + + agent_name = str(uuid.uuid4()) + if self._use_torch_transfer: + if backend != "UCX": + raise NotImplementedError( + "Experimental PyTorch transfer currently supports only " + "the NIXL UCX backend" + ) + if self.enable_staging: + raise NotImplementedError( + "Experimental PyTorch transfer does not yet support " + "SGLang staging buffers" + ) + if len(set(self.kv_args.kv_data_mem_kinds)) > 1: + raise NotImplementedError( + "Experimental PyTorch transfer does not yet support mixed " + "KV memory kinds" + ) + from sglang.srt.disaggregation.nixl.torch_transfer import ( + TorchTransferAgent, + ) + + self.agent = TorchTransferAgent( + agent_name, + owner=self, + background_progress=True, + num_threads=num_threads, + ) + else: + try: + from nixl._api import ( + nixl_agent, + nixl_agent_config, + nixl_thread_sync_t, + ) + except ImportError as e: + raise ImportError( + "Please install NIXL by following the instructions at " + "https://github.com/ai-dynamo/nixl/blob/main/README.md " + "to run SGLang with NixlTransferEngine." + ) from e + + # Bootstrap and transfer-worker threads share the agent. + agent_config = nixl_agent_config( + backends=[], + num_threads=num_threads, + sync_mode=nixl_thread_sync_t.NIXL_THREAD_SYNC_STRICT, + ) + self.agent = nixl_agent(agent_name, agent_config) run_with_deadline( lambda: self.agent.create_backend(backend, backend_params), timeout_s=envs.SGLANG_DISAGGREGATION_ENGINE_INIT_TIMEOUT.get(), @@ -471,11 +688,11 @@ def __init__( f"NIXL backend '{backend}' not found. Available: {available_plugins}. " f"Please install the required NIXL plugin or choose from: {available_plugins}" ) - logger.info(f"NIXL KVManager initialized with backend: {backend}") + data_plane = "PyTorch transfer" if self._use_torch_transfer else "native NIXL" + logger.info(f"NIXL KVManager initialized with {data_plane} backend: {backend}") self.register_buffer_to_engine() - self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() self.kv_buffer_tensors = None self.prep_handles: Dict[str, Any] = {} self.prep_handle_slice_src: Optional[Tuple[Any, int, int, int]] = ( @@ -496,6 +713,10 @@ def __init__( FastQueue() for _ in range(transfer_queue_size) ] self.exceptions: Dict[int, Exception] = {} + # Peer registration happens before a room-specific bootstrap + # message. Preserve failures by remote agent so the later room can + # transition to Failed instead of waiting after the listener dies. + self._bootstrap_errors: Dict[str, 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) @@ -518,14 +739,14 @@ def __init__( ).start() self._start_bootstrap_thread() elif self.disaggregation_mode == DisaggregationMode.DECODE: - self.transfer_statuses: Dict[int, TransferStatus] = defaultdict( - TransferStatus - ) + # Receiver construction binds each entry to an exact generation. + # Never auto-create room-only status from an untrusted notification. + self.transfer_statuses: Dict[int, TransferStatus] = {} if self.enable_staging: self._init_staging_decode_ctx() self._staging_handler = None - if self.enable_staging or self.enable_deferred_decode_kv_release: - self._start_decode_listener_thread() + # ABORT_ACK is mandatory even without heterogeneous-TP staging. + self._start_decode_listener_thread() self._start_heartbeat_checker_thread() else: raise ValueError( @@ -602,15 +823,27 @@ def _start_decode_listener_thread(self): def decode_listener_thread(): while True: msg = self.server_socket.recv_multipart() - if msg[0] == b"STAGING_REQ": + if msg[0] == STAGING_REQ_WIRE_TAG: 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")) + # Only a complete exact-generation ACK can release memory. + try: + if len(msg) != 4: + raise ValueError( + "ABORT_ACK must contain exactly four frames" + ) + room = int(msg[1].decode("ascii")) + prefill_rank = int(msg[2].decode("ascii")) + generation = validate_transfer_generation( + msg[3].decode("ascii") + ) + except (UnicodeDecodeError, ValueError): + logger.warning( + "Ignoring generation-less or malformed ABORT_ACK" ) + else: + self.note_abort_ack(room, generation, prefill_rank) continue logger.warning( "decode_listener_thread: unexpected message tag %s", @@ -703,16 +936,27 @@ def _prep_equal_tp_dlist( f"xfer_len={xfer_len}, item_len={item_len}, mem_kind={mem_kind}" ) n = num_slots if num_slots is not None else (data_len // item_len) - addrs = np.arange(n, dtype=np.uint64) * np.uint64(item_len) + base_ptr - arrays.append( - np.column_stack( - [ - addrs, - np.full(n, xfer_len, dtype=np.uint64), - np.full(n, device_id, dtype=np.uint64), - ] + if self._use_torch_transfer: + # Core indexed plans flatten the logical blocks represented by + # each strided Region. Keep one compact row per allocation; + # existing layer-major/slot-minor indices remain unchanged. + arrays.append( + np.asarray( + [[base_ptr, xfer_len, device_id, item_len, n]], + dtype=np.uint64, + ) + ) + else: + addrs = np.arange(n, dtype=np.uint64) * np.uint64(item_len) + base_ptr + arrays.append( + np.column_stack( + [ + addrs, + np.full(n, xfer_len, dtype=np.uint64), + np.full(n, device_id, dtype=np.uint64), + ] + ) ) - ) prep_handle = self.agent.prep_xfer_dlist(peer_name, np.vstack(arrays), mem_kind) assert prep_handle is not None, ( @@ -1095,13 +1339,17 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) kv_chunk: TransferKVChunk = queue.get() room = kv_chunk.room handles: List[Any] = [] + batch_active = False try: + if self._use_torch_transfer: + self.agent.begin_handle_batch(handles) + batch_active = True if room not in self.request_status: logger.debug( "Skipping chunk for room %s because it has been cleared", room, ) - self._staging_outstanding.pop(room, None) + self._maybe_ack_drained_abort(room) continue # Counted at dequeue, before the status check, so @@ -1109,14 +1357,15 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) # the predicate the abort ack relies on. The flag survives # re-enqueue on defer. if not kv_chunk.staging_counted: - self._staging_outstanding[room] += 1 + self._count_transfer_writer(room) kv_chunk.staging_counted = True if self.check_status(room) == KVPoll.Failed: - self._staging_outstanding.pop(room, None) - if self.enable_deferred_decode_kv_release: - # Skipped => nothing written for this aborted room; ack. - self._maybe_ack_drained_abort(room) + remaining_writers = self._uncount_transfer_writer(room) + kv_chunk.staging_counted = False + if remaining_writers == 0 and self.mark_transfer_quiescent(room): + self.transfer_infos.pop(room, None) + self.req_to_decode_prefix_len.pop(room, None) continue room_transfer_infos = self.transfer_infos.get(room) @@ -1126,9 +1375,25 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) "has been cleared", room, ) - self._staging_outstanding.pop(room, None) + remaining_writers = self._uncount_transfer_writer(room) + kv_chunk.staging_counted = False + if remaining_writers == 0 and self.check_status(room) in ( + KVPoll.Success, + KVPoll.Failed, + ): + self.mark_transfer_quiescent(room) continue + # Early-send can be queued while the preceding forward is still + # writing the source KV pages. Consume the event once, before + # any path can read those pages (staging, DCP packing, or direct + # NIXL/PyTorch transfer). Clear it before waiting so a staging + # defer/re-enqueue cannot repeat the host synchronization. + wait_event = kv_chunk.wait_event + kv_chunk.wait_event = None + if wait_event is not None: + wait_event.synchronize() + # Lazily build a per-worker staging strategy bound to this # worker's private staging buffer (matches mooncake). if ( @@ -1193,7 +1458,7 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) src_prefill_kv_indices = kv_chunk.prefill_kv_indices notif = ( - f"{req.room}_kv_{kv_chunk.chunk_id}" + f"{req.completion_notification_prefix}_kv_{kv_chunk.chunk_id}" f"_{int(kv_chunk.is_last_chunk)}_{self.transfer_source_rank}" ) @@ -1288,13 +1553,12 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) ), ) else: - handles.extend( - self.send_kvcache_mixed( - req.agent_name, - src_prefill_kv_indices, - chunked_dst_kv_indice, - notif, - ) + self.send_kvcache_mixed( + req.agent_name, + src_prefill_kv_indices, + chunked_dst_kv_indice, + notif, + handles, ) else: kv_xfer_handle = self.send_kvcache_slice( @@ -1304,27 +1568,26 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) notif, ) - if kv_xfer_handle is not None: + if kv_xfer_handle is not None and not self._use_torch_transfer: handles.append(kv_xfer_handle) if kv_chunk.is_last_chunk: dst_info = self.decode_kv_args_table[req.agent_name] if kv_chunk.state_indices: - state_xfer_handles = self.maybe_send_extra( + self.maybe_send_extra( req.agent_name, kv_chunk.state_indices, dst_info.dst_state_data_ptrs, req.dst_state_indices, dst_info.gpu_id, - f"{req.room}_state_{self.transfer_source_rank}", + f"{req.completion_notification_prefix}_state_" + f"{self.transfer_source_rank}", decode_tp_size, decode_tp_rank=dst_info.decode_tp_rank, dst_state_item_lens=dst_info.dst_state_item_lens, dst_state_dim_per_tensor=dst_info.dst_state_dim_per_tensor, dst_state_layer_ids=dst_info.dst_state_layer_ids, - ) - handles.extend( - h for h in state_xfer_handles if h is not None + handles=handles, ) if kv_chunk.prefill_aux_index is None: @@ -1332,13 +1595,18 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) # A no-KV notification still identifies its PP source. # Empty non-final chunks do not consume chunk IDs, so a # final no-KV chunk_id equals the prior KV chunk count. - aux_notif = f"{req.room}_aux" if ( len(kv_chunk.prefill_kv_indices) == 0 or not self.kv_args.kv_data_ptrs ): - aux_notif += ( - f"_nokv_{self.transfer_source_rank}_{kv_chunk.chunk_id}" + aux_notif = ( + f"{req.completion_notification_prefix}_aux_nokv_" + f"{self.transfer_source_rank}_{kv_chunk.chunk_id}" + ) + else: + aux_notif = ( + f"{req.completion_notification_prefix}_aux_" + f"{self.transfer_source_rank}" ) aux_xfer_handle = self.send_aux( req.agent_name, @@ -1347,17 +1615,29 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) req.dst_aux_index, aux_notif, ) - handles.append(aux_xfer_handle) + if not self._use_torch_transfer: + handles.append(aux_xfer_handle) if staging_deferred: # Chunk has been re-enqueued; do not advance status. continue while handles: + if self._use_torch_transfer: + # Advance provider progress once. The following loop is + # the sole status scan for caller-owned batch Works. + self.agent.progress() all_done = True for handle in handles: state = self.agent.check_xfer_state(handle) if state == "ERR": + deferred_error = ( + self.agent.pop_xfer_error(handle) + if self._use_torch_transfer + else None + ) + if deferred_error is not None: + raise deferred_error raise RuntimeError( f"NIXL transfer encountered ERR room={room}" ) @@ -1367,11 +1647,8 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) break time.sleep(0) - self._staging_outstanding[room] -= 1 - if self.enable_deferred_decode_kv_release: - # Handles all DONE => this room's writes landed; ack if it - # was aborted and nothing else is outstanding. - self._maybe_ack_drained_abort(room) + remaining_writers = self._uncount_transfer_writer(room) + kv_chunk.staging_counted = False if kv_chunk.is_last_chunk: self.update_status(room, KVPoll.Success) elif self.check_status(room) != KVPoll.Success: @@ -1379,18 +1656,17 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) # already concluded Success; don't regress the status. self.update_status(room, KVPoll.Transferring) - # Drop per-room state only when no chunk is still outstanding and - # the room has concluded: Success, or a Failed *last* chunk. A - # non-last Failed chunk keeps the room (more chunks may follow); a - # late chunk for an already-Failed room is skipped at loop top. - if self._staging_outstanding.get(room, 0) <= 0 and ( - self.check_status(room) == KVPoll.Success - or ( - kv_chunk.is_last_chunk - and self.check_status(room) == KVPoll.Failed + # Once a room is Failed, enqueue rejects all later chunks. When + # its exact outstanding count reaches zero, retain only the + # process-lifetime generation tombstone, not mutable room state. + if ( + remaining_writers == 0 + and ( + self.check_status(room) == KVPoll.Success + or self.check_status(room) == KVPoll.Failed ) + and self.mark_transfer_quiescent(room) ): - self._staging_outstanding.pop(room, None) self.transfer_infos.pop(room, None) self.req_to_decode_prefix_len.pop(room, None) if self.enable_staging and self._staging_ctx is not None: @@ -1399,7 +1675,61 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) for k in list(self._staging_ctx.prefetch_requested): if k[0] == room: self._staging_ctx.prefetch_requested.discard(k) - except Exception as e: + except BaseException as e: + if self._use_torch_transfer and handles: + if not isinstance(e, Exception): + # One identity publication retains every handle already + # appended to the prepublished batch, even if this worker + # is about to rethrow a process-control BaseException. + try: + self.agent.retain_handle_batch(handles) + except BaseException: + pass + while True: + try: + # Source-ownership barrier: only terminal, detached + # Works permit the scheduler to recycle this room. + self.agent.cancel_handles_and_wait(handles) + break + except BaseException as cleanup_error: + try: + if self.agent.handle_batch_is_quiescent(handles): + break + except BaseException: + pass + logger.warning( + "PyTorch cleanup for room %s was interrupted by %s; " + "retrying before publishing failure", + room, + type(cleanup_error).__qualname__, + ) + time.sleep(0) + if kv_chunk.staging_counted: + self._uncount_transfer_writer(room) + kv_chunk.staging_counted = False + if isinstance(e, Exception): + stored_error = e + try: + failure_message = str(e) + except BaseException: + failure_message = type(e).__qualname__ + else: + try: + detail = str(e) + except BaseException: + detail = "diagnostic formatting unavailable" + failure_message = f"{type(e).__qualname__}: {detail}" + # Do not retain a fatal exception traceback and its handle + # graph after every Work has become reusable. + stored_error = RuntimeError(failure_message) + self.exceptions[room] = stored_error + self.record_failure(room, failure_message) + self.update_status(room, KVPoll.Failed) + if self._transfer_writer_count( + room + ) == 0 and self.mark_transfer_quiescent(room): + self.transfer_infos.pop(room, None) + self.req_to_decode_prefix_len.pop(room, None) # Catch all exceptions to prevent silently killing this # worker thread, but still propagate via failure_exception(). if isinstance(e, _NIXL_TRANSPORT_ERRORS): @@ -1408,11 +1738,12 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) logger.exception( f"Unexpected transfer worker error for room {room}" ) - self.exceptions[room] = e - self.record_failure(room, str(e)) - self.update_status(room, KVPoll.Failed) - # No ack here on purpose: the DONE barrier bails on the first - # ERR, so siblings may still be writing; fall back to the timeout. + # Deferred-release acknowledgement now follows the terminal drain. + if not isinstance(e, Exception): + raise + finally: + if batch_active: + self.agent.end_handle_batch(handles) def register_buffer_to_engine(self): self.kv_descs = [] @@ -1480,12 +1811,94 @@ def _add_remote_peer(self, decode_kv_args: KVArgsRegisterInfo): decode_kv_args.requires_dcp_relayout = self.requires_dcp_relayout( decode_kv_args.dst_dcp_size, decode_kv_args.dst_dcp_rank ) + if self._use_torch_transfer: + if decode_kv_args.requires_dcp_relayout: + raise NotImplementedError( + "Experimental PyTorch transfer does not support DCP relayout" + ) + if decode_kv_args.decode_tp_size != self.attn_tp_size: + raise NotImplementedError( + "Experimental PyTorch transfer currently requires equal " + "prefill and decode attention TP sizes" + ) + if len(set(decode_kv_args.dst_kv_mem_kinds)) > 1: + raise NotImplementedError( + "Experimental PyTorch transfer does not support mixed remote " + "KV memory kinds" + ) if decode_kv_args.requires_dcp_relayout: self._init_dcp_pack_buffers_once(decode_kv_args.dst_dcp_size) - self.decode_kv_args_table[agent_name] = decode_kv_args - self.agent.add_remote_agent(decode_kv_args.agent_metadata) - if self.disaggregation_mode == DisaggregationMode.PREFILL: - self._prepare_payload_xfer(decode_kv_args) + if self._use_torch_transfer: + missing = object() + prior_prepared = { + name: self.prep_handles.get(name, missing) for name in ("", agent_name) + } + try: + alias = self.agent.add_remote_agent( + decode_kv_args.agent_metadata, peer_name=agent_name + ) + if alias != agent_name: + raise RuntimeError( + "PyTorch transfer imported an unexpected peer alias " + f"{alias!r} for {agent_name!r}" + ) + if self.disaggregation_mode == DisaggregationMode.PREFILL: + self._prepare_payload_xfer(decode_kv_args) + # This table is the authoritative manager-level peer commit and + # therefore publishes last. + self.decode_kv_args_table[agent_name] = decode_kv_args + except BaseException as error: + if self.decode_kv_args_table.get(agent_name) is decode_kv_args: + self.decode_kv_args_table.pop(agent_name, None) + + # Quarantine derived prepared lists before peer rollback. The + # PyTorch path is equal-TP-only, so these are its only keys. + release_prepared = getattr(self.agent, "release_dlist_handle", None) + for name, previous in prior_prepared.items(): + current = self.prep_handles.get(name, missing) + if current is previous: + continue + if previous is missing: + self.prep_handles.pop(name, None) + else: + self.prep_handles[name] = previous + if current is not missing and callable(release_prepared): + try: + release_prepared(current) + except BaseException as cleanup_error: + add_note = getattr(error, "add_note", None) + if callable(add_note): + try: + add_note( + "prepared-list rollback also raised: " + f"{type(cleanup_error).__qualname__}" + ) + except BaseException: + pass + + # Exact metadata matching prevents a rejected duplicate from + # removing some other incarnation. TorchTransferAgent marks the + # matching peer retiring before its first provider close. + try: + rollback_peer = getattr(self.agent, "rollback_remote_agent", None) + if callable(rollback_peer): + rollback_peer(agent_name, decode_kv_args.agent_metadata) + except BaseException as cleanup_error: + add_note = getattr(error, "add_note", None) + if callable(add_note): + try: + add_note( + "PyTorch peer rollback also raised: " + f"{type(cleanup_error).__qualname__}" + ) + except BaseException: + pass + raise + else: + self.decode_kv_args_table[agent_name] = decode_kv_args + self.agent.add_remote_agent(decode_kv_args.agent_metadata) + if self.disaggregation_mode == DisaggregationMode.PREFILL: + self._prepare_payload_xfer(decode_kv_args) def _send_kvcache_generic( self, @@ -1761,6 +2174,7 @@ def send_kvcache_mixed( prefill_kv_indices: npt.NDArray[np.int32], dst_kv_indices: npt.NDArray[np.int32], notif: str, + handles: list, ): info = self.decode_kv_args_table[peer_name] segments = info.kv_xfer_segments @@ -1769,7 +2183,6 @@ def send_kvcache_mixed( raise RuntimeError(f"Missing NIXL mixed KV transfer plan for {peer_name}") num_parts = len(segments) - handles = [] for part_idx, seg in enumerate(segments): num_layers = seg.end - seg.start src_indices = repeat_indices_over_layers( @@ -1792,8 +2205,8 @@ def send_kvcache_mixed( state = self.agent.transfer(xfer_handle) if state == "ERR": raise Exception("KVSender failed to post mixed prepped transfer") - handles.append(xfer_handle) - return handles + if not self._use_torch_transfer: + handles.append(xfer_handle) def send_kvcache_slice( self, @@ -2021,7 +2434,8 @@ def _do_staging_transfer( return (None, True) notif_tag = ( - f"{req.room}_stg_{kv_chunk.chunk_id}_{int(kv_chunk.is_last_chunk)}" + f"{req.completion_notification_prefix}_stg_{kv_chunk.chunk_id}" + f"_{int(kv_chunk.is_last_chunk)}" f"_{self.transfer_source_rank}_{chunk_idx}" f"_{page_start}_{num_pages}_{req.agent_name}" ) @@ -2341,6 +2755,8 @@ def maybe_send_extra( dst_state_item_lens: List[List[int]] | None = None, dst_state_dim_per_tensor: List[List[int]] | None = None, dst_state_layer_ids: List[List[int]] | None = None, + *, + handles: list, ): """Send state per hybrid component, dispatching by state_type[i].""" state_types = getattr(self.kv_args, "state_types", []) or [] @@ -2360,7 +2776,6 @@ def maybe_send_extra( dst_state_dim_per_tensor = dst_state_dim_per_tensor or [] dst_state_layer_ids = dst_state_layer_ids or [] - handles = [] for i, st in enumerate(state_types): src_indices = ( prefill_state_indices[i] if i < len(prefill_state_indices) else None @@ -2518,8 +2933,8 @@ def maybe_send_extra( f"PD Disaggregation via NIXL does NOT support {st} hybrid models yet." ) if h is not None: - handles.append(h) - return handles + if not self._use_torch_transfer: + handles.append(h) def add_transfer_request( self, @@ -2531,6 +2946,7 @@ def add_transfer_request( aux_index: Optional[int] = None, state_indices: Optional[List] = None, num_kv_tokens: Optional[int] = None, + wait_event: Optional[object] = None, ): assert self.disaggregation_mode == DisaggregationMode.PREFILL assert not is_last_chunk or (is_last_chunk and aux_index is not None) @@ -2542,7 +2958,11 @@ def add_transfer_request( if self.enable_staging: self._prefetch_staging_reqs(bootstrap_room) - if bootstrap_room not in self.transfer_infos: + if ( + bootstrap_room not in self.request_status + or self.check_status(bootstrap_room) == KVPoll.Failed + or bootstrap_room not in self.transfer_infos + ): # Dummy rank or already cleared; nothing to enqueue. return None @@ -2562,6 +2982,7 @@ def add_transfer_request( prefill_aux_index=aux_index, state_indices=state_indices, num_kv_tokens=num_kv_tokens, + wait_event=wait_event, ) ) return None @@ -2571,82 +2992,96 @@ def update_transfer_status(self): notif_map = self.agent.get_new_notifs() for peer_name, messages in notif_map.items(): for msg in messages: - # Notification tag layouts (underscore-separated): - # kv: {room}_kv_{chunk_id}_{is_last}_{pp_rank} -> 5 fields - # kvpart:{room}_kv_{chunk_id}_{is_last}_{pp_rank}_part_{i}_{n}-> 8 fields - # stg: {room}_stg_{chunk_id}_{is_last}_{pp_rank}_{chunk_idx} - # _{page_start}_{num_pages}_{agent_name} -> 9 fields - # aux: {room}_aux -> 2 fields - # state: {room}_state_{pp_rank} -> 3 fields - # maxsplit=8 keeps everything past the 8th underscore in the - # last component, so agent_name (which may itself contain - # underscores) lands intact in components[8] for the stg path. - components = msg.decode("ascii").split("_", 8) - room = int(components[0]) - tag = components[1] + try: + generation, room, tag, fields = _parse_completion_notification(msg) + except (TypeError, UnicodeError, ValueError) as exc: + logger.debug( + "Ignoring malformed NIXL completion from %s: %s", + peer_name, + exc, + ) + continue + + # This lookup never creates state. Generation comparison occurs + # before any status, staging-ring, or scatter mutation. + status = self.transfer_statuses.get(room) + if status is None or status.transfer_generation != generation: + logger.debug( + "Ignoring stale NIXL completion for room %s generation %s", + room, + generation, + ) + continue + + if tag in ("kv", "kvpart", "stg"): + source_rank = fields[2] + elif tag in ("aux", "state"): + source_rank = fields[0] + else: # The parser is exhaustive; retain a fail-closed guard. + continue + if not status.accepts_source_rank(source_rank): + logger.debug( + "Ignoring NIXL completion for unexpected source rank %s " + "in room %s generation %s", + source_rank, + room, + generation, + ) + continue + if tag == "kv": - chunk_id = int(components[2]) - is_last_chunk = bool(int(components[3])) - pp_rank = int(components[4]) if len(components) > 4 else 0 - if len(components) > 7 and components[5] == "part": - self._track_kv_part_arrival( - room, - chunk_id, - is_last_chunk, - pp_rank, - int(components[6]), - int(components[7]), - ) - else: - self._track_kv_arrival(room, chunk_id, is_last_chunk, pp_rank) + self._track_kv_arrival(status, room, *fields) + elif tag == "kvpart": + self._track_kv_part_arrival(status, room, *fields) elif tag == "stg": - self._handle_stg_notification(components, room) + self._handle_stg_notification(status, room, fields) elif tag == "aux": - # Main's "nokv" marker carries the number of earlier KV - # chunks expected from this PP rank. - self._handle_aux_notification(room, components) + self._handle_aux_notification(status, room, fields) elif tag == "state": - pp_rank = int(components[2]) if len(components) > 2 else 0 - self.transfer_statuses[room].received_state_per_pp.add(pp_rank) + pp_rank, _component_id = fields + status.received_state_per_pp.add(pp_rank) - def _handle_stg_notification(self, components, room: int): + def _handle_stg_notification( + self, status: TransferStatus, room: int, fields: Tuple[Any, ...] + ): """Handle a staging RDMA notification tag. - Format: {room}_stg_{chunk_id}_{is_last}_{pp_rank}_{chunk_idx}_{page_start}_{num_pages}_{agent_name} + Format: nixlv2_{generation}_{room}_stg_{chunk_id}_{is_last}_{pp_rank}_{chunk_idx}_{page_start}_{num_pages}_{agent_name} """ - chunk_id = int(components[2]) - is_last_chunk = bool(int(components[3])) - pp_rank = int(components[4]) - chunk_idx = int(components[5]) - page_start = int(components[6]) - num_pages = int(components[7]) - agent_name = components[8] if len(components) > 8 else "" + ( + chunk_id, + is_last_chunk, + pp_rank, + chunk_idx, + page_start, + num_pages, + agent_name, + ) = fields # Count this notif's own arrival BEFORE _track_kv_arrival, which can # conclude the transfer and record all-ranks Success. self._handle_staging_chunk_arrived( - room, chunk_idx, page_start, num_pages, agent_name + room, chunk_idx, page_start, num_pages, pp_rank ) - self._track_kv_arrival(room, chunk_id, is_last_chunk, pp_rank) + self._track_kv_arrival(status, room, chunk_id, is_last_chunk, pp_rank) - def _handle_aux_notification(self, room: int, components: List[str]): + def _handle_aux_notification( + self, status: TransferStatus, room: int, fields: Tuple[Any, ...] + ): """Handle an aux notification and trigger last scatter if staging is complete. Notification tag layouts: - aux: {room}_aux -> 2 fields - aux (nokv): {room}_aux_nokv_{pp_rank}_{expected} -> 5 fields + aux: nixlv2_{generation}_{room}_aux_{pp_rank} + aux (nokv): nixlv2_{generation}_{room}_aux_nokv_{pp_rank}_{expected} (the last chunk had no KV pages for this rank; `expected` is the number of prior KV chunks) """ - self.transfer_statuses[room].received_aux = True + pp_rank, expected = fields + status.received_aux = True # main's "nokv" marker (decode-side radix cache hit, see #19746). - if len(components) > 3 and components[2] == "nokv": - pp_rank = int(components[3]) - expected = int(components[4]) if len(components) > 4 else 0 - self.transfer_statuses[room].expected_kvs_per_pp[pp_rank] = expected - if self.transfer_statuses[room].num_pp_ranks_expected is None: - self.transfer_statuses[ - room - ].num_pp_ranks_expected = self.required_prefill_response_num_table.get( + if expected is not None: + status.expected_kvs_per_pp[pp_rank] = expected + if status.num_pp_ranks_expected is None: + status.num_pp_ranks_expected = self.required_prefill_response_num_table.get( room, 1 ) if ( @@ -2654,30 +3089,34 @@ def _handle_aux_notification(self, room: int, components: List[str]): and self._staging_handler is not None and self._staging_handler.is_staging_room(room) ): - self._maybe_submit_last_scatter(room) + self._maybe_submit_last_scatter(room, status) def _track_kv_arrival( - self, room: int, chunk_id: int, is_last_chunk: bool, pp_rank: int + self, + status: TransferStatus, + room: int, + chunk_id: int, + is_last_chunk: bool, + pp_rank: int, ): """Update transfer status tracking for a kv chunk arrival.""" - self.transfer_statuses[room].received_kvs_per_pp[pp_rank].add(chunk_id) + status.received_kvs_per_pp[pp_rank].add(chunk_id) if is_last_chunk: - self.transfer_statuses[room].expected_kvs_per_pp[pp_rank] = chunk_id + 1 - if self.transfer_statuses[room].num_pp_ranks_expected is None: - self.transfer_statuses[ - room - ].num_pp_ranks_expected = self.required_prefill_response_num_table.get( - room, 1 + status.expected_kvs_per_pp[pp_rank] = chunk_id + 1 + if status.num_pp_ranks_expected is None: + status.num_pp_ranks_expected = ( + self.required_prefill_response_num_table.get(room, 1) ) if ( self.enable_staging and self._staging_handler is not None and self._staging_handler.is_staging_room(room) ): - self._maybe_submit_last_scatter(room) + self._maybe_submit_last_scatter(room, status) def _track_kv_part_arrival( self, + status: TransferStatus, room: int, chunk_id: int, is_last_chunk: bool, @@ -2687,7 +3126,7 @@ def _track_kv_part_arrival( ): """Track one segment of a mixed-memory KV transfer.""" if num_parts <= 1: - self._track_kv_arrival(room, chunk_id, is_last_chunk, pp_rank) + self._track_kv_arrival(status, room, chunk_id, is_last_chunk, pp_rank) return if part_idx < 0 or part_idx >= num_parts: raise RuntimeError( @@ -2697,7 +3136,6 @@ def _track_kv_part_arrival( ) key = (pp_rank, chunk_id) - status = self.transfer_statuses[room] if status.received_kv_parts_per_pp is None: status.received_kv_parts_per_pp = defaultdict(set) if status.expected_kv_parts_per_pp is None: @@ -2713,7 +3151,7 @@ def _track_kv_part_arrival( if len(parts) == num_parts: status.received_kv_parts_per_pp.pop(key, None) status.expected_kv_parts_per_pp.pop(key, None) - self._track_kv_arrival(room, chunk_id, is_last_chunk, pp_rank) + self._track_kv_arrival(status, room, chunk_id, is_last_chunk, pp_rank) def _handle_staging_chunk_arrived( self, @@ -2721,7 +3159,7 @@ def _handle_staging_chunk_arrived( chunk_idx: int, page_start: int, num_pages: int, - agent_name: str, + source_rank: int, ): """Process a staging chunk arrival via RDMA notification.""" handler = self._staging_handler @@ -2732,14 +3170,19 @@ def _handle_staging_chunk_arrived( chunk_idx, page_start, num_pages, - agent_name, + source_rank, ) - def _maybe_submit_last_scatter(self, room: int): + def _maybe_submit_last_scatter( + self, room: int, status: Optional[TransferStatus] = None + ): """Check if all kv+aux transfers are done and submit last scatter if so.""" - status = self.transfer_statuses.get(room) - if status is None: + current_status = self.transfer_statuses.get(room) + if current_status is None or ( + status is not None and current_status is not status + ): return + status = current_status if not status.received_aux: return if status.num_pp_ranks_expected is None: @@ -2763,15 +3206,22 @@ def _handle_abort_notification(self, msg: List[bytes]) -> bool: return False try: + if len(msg) != 5: + raise ValueError("ABORT must contain exactly five frames") 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 + decode_ip = msg[2].decode("ascii") + decode_port = int(msg[3].decode("ascii")) + transfer_generation = validate_transfer_generation(msg[4].decode("ascii")) except Exception as e: logger.debug(f"Ignoring malformed abort notification: {e}") return True + generation_active = ( + self.active_transfer_generation(room_to_be_aborted) == transfer_generation + ) room_active = ( - room_to_be_aborted in self.request_status + generation_active + and room_to_be_aborted in self.request_status and self.check_status(room_to_be_aborted) != KVPoll.Success ) if room_active: @@ -2796,78 +3246,168 @@ def _handle_abort_notification(self, msg: List[bytes]) -> bool: # 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) + if decode_port is not None and self.register_deferred_ack_target( + room_to_be_aborted, + transfer_generation, + decode_ip, + decode_port, + ): + self._maybe_ack_drained_abort(room_to_be_aborted, transfer_generation) + elif not generation_active: + logger.debug( + "Ignoring stale ABORT for room %s generation %s", + room_to_be_aborted, + transfer_generation, + ) return True - def _start_bootstrap_thread(self): - def bootstrap_thread(): - """This thread recvs transfer info from the decode engine""" - while True: - waiting_req_bytes = self.server_socket.recv_multipart() - logger.debug( - f"Received multipart with total byte size {sum(len(x) for x in waiting_req_bytes)}" - ) + def _handle_bootstrap_message(self, waiting_req_bytes): + """Handle one bootstrap message and persist failures for propagation.""" - # Staging: decode reports consumption watermark back to prefill - if waiting_req_bytes[0] == b"WATERMARK": - if self.enable_staging: - handle_watermark_msg(self._staging_ctx, waiting_req_bytes) - continue + agent_name: Optional[str] = None + room_id: Optional[int] = None + is_room_message = False + room_identity_authenticated = False + try: + logger.debug( + "Received multipart with total byte size %d", + sum(len(x) for x in waiting_req_bytes), + ) + if not waiting_req_bytes: + raise ValueError("NIXL bootstrap message is empty") - # Staging: decode replies with allocated staging offset - if waiting_req_bytes[0] == b"STAGING_RSP": - if self.enable_staging: - handle_staging_rsp(waiting_req_bytes, self.transfer_infos) - continue + # Staging: decode reports consumption watermark back to prefill. + if waiting_req_bytes[0] == b"WATERMARK": + if self.enable_staging: + handle_watermark_msg(self._staging_ctx, waiting_req_bytes) + return - if self._handle_abort_notification(waiting_req_bytes): - continue + # Staging: decode replies with allocated staging offset. + if waiting_req_bytes[0] == STAGING_RSP_WIRE_TAG: + if self.enable_staging: + handle_staging_rsp(waiting_req_bytes, self.transfer_infos) + return + + if self._handle_abort_notification(waiting_req_bytes): + return - assert waiting_req_bytes[0] == GUARD, ( - f"First message should be {GUARD}. Foreign traffic?" + if waiting_req_bytes[0] != GUARD: + raise ValueError(f"First message should be {GUARD}. Foreign traffic?") + payload = waiting_req_bytes[1:] + room_text = payload[0].decode("ascii") + agent_name = payload[3].decode("ascii") + if room_text == "None": + # Registration failures are agent-scoped because no room is + # available yet. A subsequent room message propagates them. + self._add_remote_peer(KVArgsRegisterInfo.from_zmq(payload)) + self._bootstrap_errors.pop(agent_name, None) + logger.debug("Register KVArgs from %s successfully", agent_name) + return + + is_room_message = True + room_id = int(room_text) + try: + incoming_generation = validate_transfer_generation( + payload[10].decode("ascii") if len(payload) > 10 else None ) - waiting_req_bytes = waiting_req_bytes[1:] - room = waiting_req_bytes[0].decode("ascii") - agent_name = waiting_req_bytes[3].decode("ascii") - if room == "None": - # Register new peer and save KV base pointers. - self._add_remote_peer( - KVArgsRegisterInfo.from_zmq(waiting_req_bytes) - ) - logger.debug(f"Register KVArgs from {agent_name} successfully") - continue - room = int(room) - if room not in self.transfer_infos: - self.transfer_infos[room] = {} - self.transfer_infos[room][agent_name] = TransferInfo.from_zmq( - waiting_req_bytes + except (AttributeError, UnicodeError, ValueError) as error: + # A room number without an exact generation cannot own or fail + # that recyclable room, nor poison this agent's later traffic. + logger.warning( + "Dropping unauthenticated NIXL metadata for room %s: %s", + room_id, + error, + ) + return + prior_error = self._bootstrap_errors.get(agent_name) + if prior_error is not None: + # Canonical generation plus the failed peer registration is the + # ownership evidence; no active-room table was necessarily + # initialized when peer setup failed. + room_identity_authenticated = True + raise prior_error + active_generation = self.active_transfer_generation(room_id) + if ( + active_generation is not None + and active_generation != incoming_generation + ): + logger.warning( + "Dropping late NIXL metadata for room %s generation %s; " + "active generation is %s", + room_id, + incoming_generation, + active_generation, + ) + return + room_identity_authenticated = active_generation == incoming_generation + transfer_info = TransferInfo.from_zmq(payload) + room_identity_authenticated = True + if not self.register_transfer_generation( + room_id, transfer_info.transfer_generation + ): + logger.warning( + "Dropping late NIXL metadata for room %s generation %s; " + "active generation is %s", + room_id, + transfer_info.transfer_generation, + self.active_transfer_generation(room_id), + ) + return + if room_id not in self.transfer_infos: + self.transfer_infos[room_id] = {} + self.transfer_infos[room_id][agent_name] = transfer_info + required_dst_info_num = self.transfer_infos[room_id][ + agent_name + ].required_dst_info_num + logger.debug( + "got info room=%s agent_name=%s required_dst_info_num=%s", + room_id, + agent_name, + required_dst_info_num, + ) + if len(self.transfer_infos[room_id]) == required_dst_info_num: + self.resolve_kv_replica_factor(self.transfer_infos[room_id]) + self.req_to_decode_prefix_len[room_id] = next( + ( + info.decode_prefix_len + for info in self.transfer_infos[room_id].values() + if info.decode_prefix_len is not None + ), + 0, + ) + logger.debug("room=%s is bootstrapped", room_id) + self.update_status(room_id, KVPoll.WaitingForInput) + except Exception as error: + if is_room_message and not room_identity_authenticated: + logger.warning( + "Dropping malformed unauthenticated NIXL room metadata: %s", + error, ) - required_dst_info_num = self.transfer_infos[room][ - agent_name - ].required_dst_info_num - logger.debug(f"got info {room=} {agent_name=} {required_dst_info_num=}") - if len(self.transfer_infos[room]) == required_dst_info_num: - self.resolve_kv_replica_factor(self.transfer_infos[room]) - self.req_to_decode_prefix_len[room] = next( - ( - info.decode_prefix_len - for info in self.transfer_infos[room].values() - if info.decode_prefix_len is not None - ), - 0, + return + if agent_name is not None: + self._bootstrap_errors.setdefault(agent_name, error) + error = self._bootstrap_errors[agent_name] + if room_id is not None: + self.exceptions[room_id] = error + self.record_failure(room_id, str(error)) + self.update_status(room_id, KVPoll.Failed) + raise + + def _start_bootstrap_thread(self): + + def bootstrap_thread(): + """This thread recvs transfer info from the decode engine""" + while True: + waiting_req_bytes = self.server_socket.recv_multipart() + try: + self._handle_bootstrap_message(waiting_req_bytes) + except Exception: + logger.exception( + "NIXL bootstrap message failed; listener remains active" ) - logger.debug(f"{room=} is bootstrapped") - self.update_status(room, KVPoll.WaitingForInput) - threading.Thread(target=bootstrap_thread).start() + threading.Thread(target=bootstrap_thread, daemon=True).start() class NixlKVSender(CommonKVSender): @@ -2894,6 +3434,16 @@ def __init__( self._send_failed = False self._send_error: Optional[Exception] = None self._transfer_start_time: Optional[float] = None + self._early_send_wait_event: Optional[object] = None + + @property + def supports_overlap_early_send(self) -> bool: + """Whether the data plane can order RDMA after a CUDA readiness event.""" + + # Both native NIXL and the PyTorch transfer provider use the common + # transfer worker, which consumes the readiness event before touching + # source memory. + return True def send( self, @@ -2915,6 +3465,9 @@ def send( ): self._transfer_start_time = time.perf_counter() + wait_event = getattr(self, "_early_send_wait_event", None) + self._early_send_wait_event = None + self.kv_mgr.add_transfer_request( self.bootstrap_room, kv_indices, @@ -2924,6 +3477,7 @@ def send( self.aux_index, state_indices, num_kv_tokens, + wait_event=wait_event, ) self._record_transfer_indices(kv_indices, state_indices) self.chunk_id += 1 @@ -2993,11 +3547,51 @@ def __init__( mgr: NixlKVManager, bootstrap_addr: str, bootstrap_room: Optional[int] = None, + transfer_generation: Optional[str] = None, ): self.started_transfer = False - super().__init__(mgr, bootstrap_addr, bootstrap_room) + super().__init__( + mgr, + bootstrap_addr, + bootstrap_room, + transfer_generation=transfer_generation, + ) + # Arm the exact expected generation before metadata can cause any + # completion. Reuse replaces old room status; late old-generation + # notifications are rejected by update_transfer_status(). + self.kv_mgr.transfer_statuses[self.bootstrap_room] = TransferStatus( + transfer_generation=self.transfer_generation + ) self.init_time = None + def init(self, prefill_dp_rank: int): + super().init(prefill_dp_rank) + if self.conclude_state == KVPoll.Failed: + return + expected_source_ranks = frozenset( + bootstrap_info["prefill_unique_rank"] + for bootstrap_info in self.bootstrap_infos + if not bootstrap_info["is_dummy"] + ) + status = self.kv_mgr.transfer_statuses.get(self.bootstrap_room) + if status is None or status.transfer_generation != self.transfer_generation: + raise RuntimeError( + "NIXL completion status lost its receiver lease identity" + ) + status.bind_expected_source_ranks(expected_source_ranks) + + def _clear_own_transfer_status(self) -> None: + status = self.kv_mgr.transfer_statuses.get(self.bootstrap_room) + if ( + status is not None + and status.transfer_generation == self.transfer_generation + ): + self.kv_mgr.transfer_statuses.pop(self.bootstrap_room, None) + + def clear(self) -> None: + self._clear_own_transfer_status() + super().clear() + def send_metadata( self, kv_indices: npt.NDArray[np.int32], @@ -3053,6 +3647,7 @@ def send_metadata( packed_state_indices, str(decode_prefix_len or 0).encode("ascii"), str(int(is_dummy)).encode("ascii"), + self.transfer_generation.encode("ascii"), ] ) except zmq.ZMQError: @@ -3085,18 +3680,18 @@ def poll(self) -> KVPoll: if not self.started_transfer: return status - # Drain notifications before enforcing the waiting deadline. The decode - # agent has no NIXL progress thread (num_threads=0), so incoming - # completion notifications are only ingested here via - # update_transfer_status(); a completion queued by NIXL at/after the - # deadline would otherwise lose to the timeout purely by poll ordering. + # Drain notifications before enforcing the waiting deadline. Both roles + # use transport background progress, but application notifications are + # ingested into room state only here via update_transfer_status(); a + # completion queued by NIXL at/after the deadline would otherwise lose + # to the timeout purely by poll ordering. self.kv_mgr.update_transfer_status() if self.kv_mgr.check_transfer_done(self.bootstrap_room): # type: ignore self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].discard( self.bootstrap_room ) self.conclude_state = KVPoll.Success - del self.kv_mgr.transfer_statuses[self.bootstrap_room] + self._clear_own_transfer_status() return self.conclude_state # type: ignore timeout_result = self._check_waiting_timeout() diff --git a/python/sglang/srt/disaggregation/nixl/torch_transfer.py b/python/sglang/srt/disaggregation/nixl/torch_transfer.py new file mode 100644 index 000000000000..7693b7ba2e97 --- /dev/null +++ b/python/sglang/srt/disaggregation/nixl/torch_transfer.py @@ -0,0 +1,2169 @@ +"""Experimental PyTorch endpoint-transfer adapter for SGLang NIXL PD.""" + +# Process-control exceptions are deliberately part of the ownership protocol: +# every BaseException boundary below reconciles or durably retains native state. +# ruff: noqa: BLE001 + +from __future__ import annotations + +import importlib +import itertools +import logging +import operator +import threading +import time +import uuid +from collections import defaultdict +from collections.abc import Callable, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Any, Literal + +logger = logging.getLogger(__name__) + +_UNKNOWN_NOTIFICATION_SOURCE = "" +_RETAINED_WORK_SWEEP_INTERVAL_S = 0.01 +_PENDING_STATES = frozenset({"CREATED", "PENDING", "RUNNING", "SUBMITTED"}) +_DONE_STATES = frozenset({"COMPLETE", "COMPLETED", "DONE", "SUCCEEDED"}) +_FAILED_STATES = frozenset({"CANCELLED", "FAILED", "TIMED_OUT"}) +_UNKNOWN_WORK_OUTCOME = object() +_FACTORY_API_VERSION = 2 + + +def _load_transfer_api() -> Any: + try: + module = importlib.import_module("torch.distributed._transfer") + except ImportError as exc: + raise ImportError( + "PyTorch endpoint transfer requires torch.distributed._transfer" + ) from exc + factory_version = getattr(module, "BACKEND_FACTORY_API_VERSION", None) + if ( + type(factory_version) is not int + or factory_version < _FACTORY_API_VERSION + or not hasattr(module, "BackendFactoryV2") + or not hasattr(module, "Endpoint") + or not hasattr(module, "RawSpan") + or not hasattr(module, "TransferOp") + or not callable( + getattr( + getattr(module, "TransferPlan", None), + "submit_indices_and_poll", + None, + ) + ) + ): + raise ImportError( + "PyTorch endpoint transfer requires FactoryV2 Core, Endpoint, RawSpan, " + "TransferOp, and allocation-minimal exact-state indexed submission " + "support" + ) + return module + + +def _register_nixl_backend() -> None: + """Load the external provider only after explicit SGLang opt-in.""" + try: + module = importlib.import_module("nixl.torch_transfer") + except ImportError as exc: + raise ImportError( + "The PyTorch NIXL transfer backend requires nixl.torch_transfer" + ) from exc + register = getattr(module, "register_torch_backend", None) + factory_version = getattr(module, "TORCH_TRANSFER_FACTORY_API_VERSION", None) + if type(factory_version) is not int or factory_version != _FACTORY_API_VERSION: + raise ImportError( + "nixl.torch_transfer must register the exact FactoryV2 provider ABI" + ) + if not callable(register): + raise ImportError( # noqa: TRY004 + "nixl.torch_transfer does not expose register_torch_backend()" + ) + register() + + +@dataclass(frozen=True) +class _RawDescriptor: + address: int + nbytes: int + device_id: int + stride: int | None = None + count: int | None = None + + @property + def extent(self) -> int: + if self.stride is None or self.count is None: + return self.nbytes + return (self.count - 1) * self.stride + self.nbytes + + +@dataclass(frozen=True) +class _RawDescriptorList: + descriptors: tuple[_RawDescriptor, ...] + memory_type: str + + +@dataclass(frozen=True) +class _PreparedList: + regions: tuple[object, ...] + remote: bool + peer_name: str | None = None + + +@dataclass(frozen=True) +class _CachedPlan: + plan: object + peer_name: str + prepared_ids: frozenset[int] + + +@dataclass +class _RegistrationRecord: + address: int + nbytes: int + device_id: int + memory_type: str + registration: object + + +def _freeze_index_vector(indices: Sequence[int]) -> Sequence[int]: + """Take immutable ownership without copying SGLang's native-int32 fast path. + + SGLang creates fresh, owning NumPy int32 arrays immediately before + ``make_prepped_xfer``. Mark those arrays read-only and retain them until + Core snapshots them during the adjacent ``transfer`` call. Non-owning + buffers are copied before freezing, and general Python sequences become an + immutable normalized tuple. Thus no caller-visible mutation can alter an + admitted request while the production fast path adds no index memcpy. + """ + + try: + view = memoryview(indices) + except TypeError: + return tuple(operator.index(value) for value in indices) + try: + native_int32 = ( + view.ndim == 1 + and view.c_contiguous + and view.itemsize == 4 + and view.format in {"i", "@i", "=i"} + ) + if not native_int32: + return tuple(operator.index(value) for value in indices) + if view.readonly: + return indices + flags = getattr(indices, "flags", None) + setflags = getattr(indices, "setflags", None) + if ( + flags is not None + and bool(getattr(flags, "owndata", False)) + and callable(setflags) + ): + setflags(write=False) + return indices + copy = getattr(indices, "copy", None) + if callable(copy): + frozen = copy() + freeze = getattr(frozen, "setflags", None) + if callable(freeze): + freeze(write=False) + return frozen + return tuple(operator.index(value) for value in indices) + finally: + view.release() + + +@dataclass(eq=False) +class _WorkHandle: + work: object | None = None + owned_plan: object | None = None + terminal_state: str | None = None + first_state: object | None = None + operation: Literal["READ", "WRITE"] | None = None + plan: object | None = None + local_indices: Sequence[int] | None = None + remote_indices: Sequence[int] | None = None + notification: bytes | None = None + peer_name: str | None = None + cache_key: tuple[str, int, int] | None = None + # The production token is the already allocated handle itself. This avoids + # a second per-attempt allocation; Core Work.close breaks the temporary + # Work-to-handle edge before the adapter drops its registry ownership. + submission_interrupted: bool = False + _state_lock: threading.Lock = field( + default_factory=threading.Lock, init=False, repr=False + ) + _status_in_progress: bool = field(default=False, init=False, repr=False) + _recovery_in_progress: bool = field(default=False, init=False, repr=False) + _release_in_progress: bool = field(default=False, init=False, repr=False) + _native_submit_entered: bool = field(default=False, init=False, repr=False) + phase: Literal[ + "prepared", + "posting", + "recovery_pending", + "posted", + "terminal", + "released", + ] = "prepared" + + +class TorchTransferAgent: + """NIXL-agent-shaped boundary around ``torch.distributed._transfer``. + + It exists to validate the current SGLang data plane without spreading an + experimental PyTorch API through ``conn.py``. Core ``MULTIPLE`` mode lets + independent transfer-worker calls overlap while a short lifecycle gate + drains every admitted call before close. Submission + deliberately uses Core's ``caller_ready`` mode: the common transfer worker + consumes the early-send CUDA event before staging, packing, or direct source + access, and it waits for terminal Work before any chunk is reusable. + """ + + def __init__( + self, + name: str, + *, + owner: object, + background_progress: bool, + num_threads: int = 0, + endpoint_backend: str = "nixl", + ) -> None: + self.name = name + self._owner = owner + self._background_progress = background_progress + self._num_threads = num_threads + self._endpoint_backend = endpoint_backend + self._api = _load_transfer_api() + self._core_transfer_operations = { + "READ": self._api.TransferOp.READ, + "WRITE": self._api.TransferOp.WRITE, + } + self._work_state_outcomes = self._complete_work_state_outcomes(self._api) + self._lock = threading.RLock() + self._lifecycle_cv = threading.Condition(self._lock) + self._admission_local = threading.local() + self._active_calls = 0 + self._closing = False + self._close_owner: int | None = None + self._endpoint: Any = None + self._endpoint_recovery: Any = None + self._transport_backend: str | None = None + self._send_notification_sync: Callable[[object, bytes], None] | None = None + self._registrations: list[_RegistrationRecord] = [] + self._registration_groups: list[tuple[object, ...]] = [] + self._registration_ids = itertools.count() + self._peers: dict[str, object] = {} + self._peer_wire_metadata: dict[str, bytes] = {} + self._source_aliases: dict[object, str] = {} + self._plan_cache: dict[tuple[str, int, int], _CachedPlan] = {} + self._closing_plans: dict[tuple[str, int, int], _CachedPlan] = {} + self._retiring_peers: set[str] = set() + self._works: dict[int, _WorkHandle] = {} + self._batch_handles = threading.local() + self._active_handle_batches: dict[int, list[_WorkHandle]] = {} + self._retained_handle_batches: dict[int, list[_WorkHandle]] = {} + self._detached_works: list[_WorkHandle] = [] + self._retained_work_ids: set[int] = set() + self._sweep_wakeup = threading.Event() + self._sweep_thread: threading.Thread | None = None + # Set before entering any Core child constructor and cleared only after + # authoritative adapter publication or a complete identity reconcile. + self._construction_dirty = False + + def create_backend(self, backend: str, backend_params: Mapping[str, str]) -> None: + with self._lock: + self._retry_endpoint_cleanup_unlocked() + if self._endpoint is not None: + raise RuntimeError("PyTorch transfer endpoint is already initialized") + options = { + "backends": [backend], + "backend_init_params": {backend: dict(backend_params)}, + "transfer_backends": [backend], + "num_threads": self._num_threads, + } + if self._endpoint_backend == "nixl": + _register_nixl_backend() + endpoint_id = str(uuid.uuid4()) + progress_mode = ( + self._api.ProgressMode.BACKGROUND + if self._background_progress + else self._api.ProgressMode.MANUAL + ) + thread_mode = self._api.ThreadMode.MULTIPLE + endpoint = None + try: + endpoint = self._api.Endpoint( + self.name, + backend=self._endpoint_backend, + endpoint_id=endpoint_id, + progress_mode=progress_mode, + # Core/provider synchronization owns native concurrency. + # Adapter locks protect only short Python state transitions. + thread_mode=thread_mode, + options=options, + ) + required_capabilities = ( + "raw_spans", + "strided_regions", + "write", + "attached_notifications", + "standalone_notifications", + ) + missing = [ + name + for name in required_capabilities + if not self._supports_capability(endpoint, name) + ] + if missing: + raise NotImplementedError( + "The selected PyTorch transfer backend is missing required " + f"capabilities: {', '.join(missing)}" + ) + send_notification_sync = getattr( + endpoint, "send_notification_sync", None + ) + if not ( + self._supports_capability(endpoint, "synchronous_notification_send") + and callable(send_notification_sync) + ): + send_notification_sync = None + # Derived publication precedes the authoritative Endpoint owner. + self._send_notification_sync = send_notification_sync + self._transport_backend = backend + self._endpoint = endpoint + except BaseException as error: + # A process-control exception may land after the authoritative + # assignment but before this method returns. The failed call + # still owns rollback: never leave a usable Endpoint published. + if endpoint is not None and self._endpoint is endpoint: + self._endpoint = None + self._send_notification_sync = None + self._transport_backend = None + + candidate = endpoint + if candidate is None: + try: + recovery_endpoint = getattr(error, "recovery_endpoint", None) + except BaseException as lookup_error: + recovery_endpoint = None + self._add_exception_note( + error, + "reading recovery_endpoint also raised: " + f"{type(lookup_error).__qualname__}", + ) + if recovery_endpoint is not None: + try: + trusted = self._matches_recovery_endpoint( + recovery_endpoint, + endpoint_id=endpoint_id, + progress_mode=progress_mode, + thread_mode=thread_mode, + ) + except BaseException as match_error: + trusted = False + self._add_exception_note( + error, + "validating recovery_endpoint also raised: " + f"{type(match_error).__qualname__}", + ) + if trusted: + candidate = recovery_endpoint + else: + # Exact identity fields are the ownership proof. + # Even another real Core Endpoint can belong to an + # unrelated live caller and must not be mutated. + self._add_exception_note( + error, + "ignored a non-matching recovery_endpoint during " + "SGLang endpoint construction", + ) + if candidate is not None: + self._endpoint_recovery = candidate + try: + self._retry_endpoint_cleanup_unlocked() + except BaseException as cleanup_error: + self._add_exception_note( + error, + "endpoint construction cleanup also raised: " + f"{type(cleanup_error).__qualname__}", + ) + raise + + def get_plugin_list(self) -> list[str]: + with self._lock: + return [self._transport_backend] if self._transport_backend else [] + + def register_memory( + self, descriptors: Sequence[Sequence[object]], memory_type: str + ) -> tuple[object, ...]: + with self._lock: + self._reconcile_construction_state_unlocked() + endpoint = self._require_endpoint() + self._construction_dirty = True + registrations = [] + records = [] + group = None + try: + for raw in self._normalize_descriptors(descriptors): + span = self._api.RawSpan( + address=raw.address, + nbytes=raw.extent, + device=self._device(memory_type, raw.device_id), + memory_type=memory_type, + owner=self._owner, + ) + name = f"sglang_region_{next(self._registration_ids)}" + registration = endpoint.register(span, name=name) + registrations.append(registration) + records.append( + _RegistrationRecord( + address=raw.address, + nbytes=raw.extent, + device_id=raw.device_id, + memory_type=memory_type, + registration=registration, + ) + ) + group = tuple(registrations) + # The complete registration graph is authoritative only after + # both indexes publish. Exception recovery removes either partial. + self._registration_groups.append(group) + self._registrations.extend(records) + self._construction_dirty = False + return group + except BaseException as error: + self._registrations = [ + record + for record in self._registrations + if not any(record is candidate for candidate in records) + ] + if group is not None: + self._registration_groups = [ + candidate + for candidate in self._registration_groups + if candidate is not group + ] + self._reconcile_after_construction_error_unlocked(error, "registration") + raise + + def deregister_memory(self, registrations: Sequence[object]) -> None: + with self._lock: + self._require_endpoint() + self._reconcile_construction_state_unlocked() + self._require_no_unsettled_interrupted_submission("deregister memory") + for registration in reversed(tuple(registrations)): + registration.close() + removed = set(map(id, registrations)) + self._registrations = [ + record + for record in self._registrations + if id(record.registration) not in removed + ] + self._registration_groups = [ + remaining + for group in self._registration_groups + if ( + remaining := tuple( + registration + for registration in group + if id(registration) not in removed + ) + ) + ] + + def get_agent_metadata(self) -> bytes: + with self._lock: + endpoint = self._require_endpoint() + metadata = endpoint.export_metadata( + [record.registration for record in self._registrations] + ) + if isinstance(metadata, bytes): + return metadata + to_bytes = getattr(metadata, "to_bytes", None) + if callable(to_bytes): + return to_bytes() + try: + return bytes(metadata) + except (TypeError, ValueError) as exc: + raise TypeError("Endpoint metadata is not bytes-compatible") from exc + + def add_remote_agent(self, metadata: bytes, peer_name: str | None = None) -> str: + with self._lock: + self._reconcile_construction_state_unlocked() + metadata = bytes(metadata) + if peer_name is not None and peer_name in self._retiring_peers: + self.remove_remote_agent(peer_name) + if peer_name is not None and peer_name in self._peers: + if self._peer_wire_metadata.get(peer_name) == metadata: + return peer_name + raise ValueError( + f"Peer {peer_name!r} is already registered with different metadata" + ) + import_options = {} if peer_name is None else {"expected_name": peer_name} + self._construction_dirty = True + peer = None + alias = None + aliases = None + previous_aliases = self._source_aliases + try: + peer = self._require_endpoint().import_peer(metadata, **import_options) + source_ids = self._peer_source_ids(peer) + source_id = source_ids[0] + alias = peer_name or source_id + if alias in self._peers: + peer.close() + raise ValueError(f"Peer {alias!r} is already registered") + if not callable(getattr(peer, "unsafe_region", None)): + peer.close() + raise NotImplementedError( + "Imported peer does not support provider-validated raw regions" + ) + aliases = dict(previous_aliases) + for source in source_ids: + aliases[source] = alias + incarnation = getattr(peer, "incarnation", None) + if incarnation is not None: + aliases[(source_id, str(incarnation))] = alias + self._source_aliases = aliases + self._peer_wire_metadata[alias] = metadata + # The exact Core peer is the authoritative adapter owner and + # commits last, after every derived identity index. + self._peers[alias] = peer + self._construction_dirty = False + return alias + except BaseException as error: + if alias is None or self._peers.get(alias) is not peer: + if aliases is not None and self._source_aliases is aliases: + self._source_aliases = previous_aliases + if alias is not None: + self._peer_wire_metadata.pop(alias, None) + self._reconcile_after_construction_error_unlocked(error, "peer") + else: + # A dict STORE may commit and then raise. Exact map identity + # proves adoption, so retry is idempotent and no scan is due. + self._construction_dirty = False + raise + + def rollback_remote_agent(self, peer_name: str, metadata: bytes) -> bool: + """Remove only the peer carrying this exact bootstrap wire metadata.""" + + with self._lock: + self._reconcile_construction_state_unlocked() + if self._peer_wire_metadata.get(peer_name) != bytes(metadata): + return False + self.remove_remote_agent(peer_name) + return True + + def remove_remote_agent(self, peer_name: str) -> None: + with self._lock: + self._require_endpoint() + self._reconcile_construction_state_unlocked() + peer = self._peers.get(peer_name) + if peer is None: + self._retiring_peers.discard(peer_name) + return + # Quarantine before the first provider close. If close partially + # commits and raises, no future submission may reuse this peer. + self._retiring_peers.add(peer_name) + self._require_no_unsettled_interrupted_submission("remove a remote agent") + self._release_cached_plans(lambda cached: cached.peer_name == peer_name) + if not getattr(peer, "closed", False): + peer.close() + self._source_aliases = { + source: alias + for source, alias in self._source_aliases.items() + if alias != peer_name + } + self._peer_wire_metadata.pop(peer_name, None) + self._peers.pop(peer_name, None) + self._retiring_peers.discard(peer_name) + + def get_xfer_descs( + self, descriptors: Sequence[Sequence[object]], memory_type: str + ) -> _RawDescriptorList: + return _RawDescriptorList(self._normalize_descriptors(descriptors), memory_type) + + def prep_xfer_dlist( + self, + peer_name: str, + descriptors: Sequence[Sequence[object]], + memory_type: str, + ) -> _PreparedList: + with self._lock: + self._require_endpoint() + raw = self._normalize_descriptors(descriptors) + if peer_name: + regions = self._resolve_remote(peer_name, raw, memory_type) + return _PreparedList(regions, remote=True, peer_name=peer_name) + return _PreparedList(self._resolve_local(raw, memory_type), remote=False) + + def release_dlist_handle(self, handle: _PreparedList) -> None: + with self._lock: + self._require_endpoint() + prepared_id = id(handle) + self._release_cached_plans( + lambda cached: prepared_id in cached.prepared_ids + ) + + def make_prepped_xfer( + self, + operation: str, + source: _PreparedList, + source_indices: Sequence[int], + destination: _PreparedList, + destination_indices: Sequence[int], + notification: bytes = b"", + ) -> _WorkHandle: + with self._admitted_call(): + operation = self._normalize_operation(operation) + local, remote, local_indices, remote_indices = self._orient_prepared( + operation, + source, + source_indices, + destination, + destination_indices, + ) + local_indices = _freeze_index_vector(local_indices) + remote_indices = _freeze_index_vector(remote_indices) + key = (operation, id(local), id(remote)) + wire_notification = bytes(notification) if notification else None + handle = None + with self._lock: + if remote.peer_name in self._retiring_peers: + raise RuntimeError( + f"Peer {remote.peer_name!r} teardown is in progress" + ) + if key in self._closing_plans: + raise RuntimeError("Prepared transfer plan teardown is in progress") + cached = self._plan_cache.get(key) + if cached is None: + assert remote.peer_name is not None + self._reconcile_construction_state_unlocked() + self._construction_dirty = True + cached = None + try: + plan = self._require_endpoint().prepare( + local=local.regions, + remote=remote.regions, + indexed=True, + ) + cached = _CachedPlan( + plan=plan, + peer_name=remote.peer_name, + prepared_ids=frozenset({id(local), id(remote)}), + ) + self._plan_cache[key] = cached + self._construction_dirty = False + except BaseException as error: + if cached is None or self._plan_cache.get(key) is not cached: + self._reconcile_after_construction_error_unlocked( + error, "indexed plan" + ) + else: + # Exact cache identity proves the STORE committed + # even when its return boundary was interrupted. + self._construction_dirty = False + raise + # Publish the plan and immutable index ownership in one state + # transition. Core snapshots and validates the indices during + # its allocation-minimal direct one-shot submit. + handle = _WorkHandle( + operation=operation, + plan=cached.plan, + local_indices=local_indices, + remote_indices=remote_indices, + notification=wire_notification, + peer_name=remote.peer_name, + cache_key=key, + ) + return self._track_work(handle) + + def initialize_xfer( + self, + operation: str, + source: _RawDescriptorList, + destination: _RawDescriptorList, + peer_name: str, + notification: bytes = b"", + ) -> _WorkHandle: + with self._lock: + self._require_endpoint() + operation = self._normalize_operation(operation) + if peer_name in self._retiring_peers: + raise RuntimeError(f"Peer {peer_name!r} teardown is in progress") + if operation == "WRITE": + local_regions = self._resolve_local( + source.descriptors, source.memory_type + ) + remote_regions = self._resolve_remote( + peer_name, destination.descriptors, destination.memory_type + ) + else: + local_regions = self._resolve_local( + destination.descriptors, destination.memory_type + ) + remote_regions = self._resolve_remote( + peer_name, source.descriptors, source.memory_type + ) + self._reconcile_construction_state_unlocked() + self._construction_dirty = True + plan = None + handle = None + try: + plan = self._require_endpoint().prepare( + local=local_regions, + remote=remote_regions, + indexed=False, + ) + handle = _WorkHandle( + owned_plan=plan, + operation=operation, + plan=plan, + notification=bytes(notification) if notification else None, + peer_name=peer_name, + ) + tracked = self._track_work(handle) + self._construction_dirty = False + return tracked + except BaseException as error: + if handle is not None and self._works.get(id(handle)) is handle: + try: + self._release_work(handle) + except BaseException as cleanup_error: + self._retained_work_ids.add(id(handle)) + self._start_retained_work_sweeper() + self._add_exception_note( + error, + "one-off handle cleanup also raised: " + f"{type(cleanup_error).__qualname__}", + ) + self._reconcile_after_construction_error_unlocked(error, "one-off plan") + raise + + def transfer(self, handle: _WorkHandle) -> str: + with self._admitted_call() as endpoint: + should_post = False + with handle._state_lock: + if handle.terminal_state is not None: + return handle.terminal_state + if handle.phase == "released": + raise RuntimeError("PyTorch transfer handle is released") + if handle.phase == "posting": + raise RuntimeError( + "PyTorch transfer submission is already in progress" + ) + if handle.phase == "prepared": + # Publish submission ownership before native entry. The lock + # is released before submit_and_poll/write/read. + handle.phase = "posting" + should_post = True + recovery_pending = handle.phase == "recovery_pending" + if recovery_pending: + self._retry_interrupted_post_recovery(handle, schedule=True) + with handle._state_lock: + recovery_pending = handle.phase == "recovery_pending" + if recovery_pending: + return "PROC" + if should_post: + try: + self._post(handle, endpoint) + except BaseException as error: + # Recover provider ownership before any diagnostic formatting: + # an exception may itself have hostile __str__/repr hooks. + with handle._state_lock: + native_submit_entered = handle._native_submit_entered + if native_submit_entered: + self._handle_interrupted_submission(handle, error) + else: + # Validation/teardown rejection precedes provider entry, + # so there is no hidden Work to reconcile. + with handle._state_lock: + handle.terminal_state = "ERR" + handle.phase = "terminal" + self._release_work(handle) + with self._lock: + retained = id(handle) in self._works + with handle._state_lock: + retained_phase = handle.phase + if retained and retained_phase in ( + "posted", + "recovery_pending", + ): + # Preserve immutable diagnostics only. Retaining the + # exception would retain this stack through a long DMA. + try: + details = (type(error).__qualname__, str(error)) + except BaseException: + details = ( + "BaseException", + "diagnostic formatting unavailable", + ) + try: + handle.__dict__["_submission_error"] = details + except BaseException: # noqa: S110 + pass + # Return the already tracked handle to conn.py's source + # lifetime barrier. The error is surfaced after terminality. + return "PROC" + raise + state = self._work_state(handle, release=False) + if state == "ERR": + # Current conn.py callers raise immediately on a post failure + # and never issue the later release call. Terminal failure is + # already authoritative, so release it here; a failed cleanup + # remains tracked and is retried by the retained-work sweeper. + try: + self._release_work(handle) + except Exception as exc: + logger.warning("PyTorch failed Work cleanup failed: %s", exc) + with self._lock: + self._retained_work_ids.add(id(handle)) + self._start_retained_work_sweeper() + return state + + def check_xfer_state(self, handle: _WorkHandle) -> str: + with self._admitted_call(): + return self._work_state(handle, release=True) + + def begin_handle_batch(self, handles: list[_WorkHandle]) -> None: + """Publish a worker-owned sink before that worker can create a Work.""" + with self._lock: + self._require_endpoint() + if getattr(self._batch_handles, "value", None) is not None: + raise RuntimeError("PyTorch transfer handle batch is already active") + batch_id = id(handles) + existing = self._active_handle_batches.get(batch_id) + if existing is not None and existing is not handles: + raise RuntimeError("PyTorch transfer handle batch identity collision") + # Publish the list itself before any Core call. Every later handle + # append is therefore already reachable if its worker is interrupted. + try: + self._active_handle_batches[batch_id] = handles + self._batch_handles.value = handles + except BaseException: + # conn.py has not observed a successful begin and therefore + # cannot own cleanup yet. Roll back either store even when the + # store itself committed before raising. + if getattr(self._batch_handles, "value", None) is handles: + del self._batch_handles.value + if self._active_handle_batches.get(batch_id) is handles: + self._active_handle_batches.pop(batch_id, None) + raise + + def end_handle_batch(self, handles: list[_WorkHandle]) -> None: + with self._lock: + self._require_endpoint() + current = getattr(self._batch_handles, "value", None) + if current is not handles: + raise RuntimeError("PyTorch transfer handle batch ownership mismatch") + del self._batch_handles.value + self._active_handle_batches.pop(id(handles), None) + self._prune_retained_handle_batches_unlocked() + if id(handles) in self._retained_handle_batches: + self._start_retained_work_sweeper() + + def retain_handle_batch(self, handles: list[_WorkHandle]) -> None: + """Durably retain a whole worker batch with one identity publication.""" + + with self._lock: + self._require_endpoint() + active = self._active_handle_batches.get(id(handles)) + if active is not None and active is not handles: + raise RuntimeError("PyTorch transfer handle batch identity collision") + self._retained_handle_batches[id(handles)] = handles + self._prune_retained_handle_batches_unlocked() + if id(handles) in self._retained_handle_batches: + self._start_retained_work_sweeper() + + def handle_batch_is_quiescent(self, handles: Sequence[_WorkHandle]) -> bool: + with self._lock: + return not any(self._works.get(id(handle)) is handle for handle in handles) + + @staticmethod + def pop_xfer_error(handle: _WorkHandle) -> RuntimeError | None: + """Return deferred lost-return diagnostics after terminal settlement.""" + + details = handle.__dict__.pop("_submission_error", None) + if details is None: + return None + error_type, message = details + return RuntimeError( + "PyTorch transfer submission return was interrupted by " + f"{error_type}: {message}" + ) + + def progress(self) -> None: + """Advance the provider and reap detached notification requests. + + SGLang calls this before checking one batch of transfer handles. Regular + Works are intentionally not scanned here: the following per-handle loop + is their single status pass. Detached notification Works have no caller + to check them, so this method still owns their terminal cleanup. + """ + + # Optimistic empty fast path: a concurrent sender publishes the Work in + # ``_works`` before appending here, so missing that append only defers + # detached cleanup to the next poll/close and cannot lose ownership. + if self._background_progress and not self._detached_works: + return + with self._admitted_call(): + self._advance_provider() + self._reap_detached_works() + + def cancel_handles(self, handles: Sequence[_WorkHandle]) -> None: + """Best-effort cancel siblings while retaining non-terminal Works.""" + + with self._admitted_call(): + for handle in handles: + with self._lock: + tracked = id(handle) in self._works + if not tracked: + continue + with handle._state_lock: + phase = handle.phase + if phase == "prepared": + handle.terminal_state = "ERR" + handle.phase = "terminal" + if phase == "prepared": + try: + self._release_work(handle) + except Exception as exc: + logger.warning( + "PyTorch prepared transfer cleanup failed: %s", exc + ) + with self._lock: + self._retained_work_ids.add(id(handle)) + self._start_retained_work_sweeper() + continue + if phase == "posting": + with self._lock: + self._retained_work_ids.add(id(handle)) + continue + if phase == "recovery_pending": + self._retry_interrupted_post_recovery(handle, schedule=True) + with handle._state_lock: + recovery_pending = handle.phase == "recovery_pending" + if recovery_pending: + continue + if self._work_state(handle, release=False) in ("DONE", "ERR"): + continue + try: + with handle._state_lock: + work = handle.work + assert work is not None + work.cancel() + except Exception as exc: + logger.warning("PyTorch transfer cancellation failed: %s", exc) + self._progress() + with self._lock: + tracked_handles = tuple( + handle + for handle in handles + if self._works.get(id(handle)) is handle + ) + retained = set() + for handle in tracked_handles: + with handle._state_lock: + recovery_pending = handle.phase == "recovery_pending" + if recovery_pending or self._work_state(handle, release=False) not in ( + "DONE", + "ERR", + ): + retained.add(id(handle)) + if retained: + with self._lock: + self._retained_work_ids.update(retained) + self._start_retained_work_sweeper() + + def cancel_handles_and_wait(self, handles: Sequence[_WorkHandle]) -> None: + """Cancel a failed batch and return only when every Work is reusable. + + This deliberately unbounded routine is used only on failure. Publishing + a room failure while an uncancellable Work is active would let SGLang + recycle KV pages that DMA may still be reading. + """ + + if isinstance(handles, list): + self.retain_handle_batch(handles) + handles = tuple(handles) + with self._admitted_call(): + self.cancel_handles(handles) + while True: + pending = False + self._advance_provider() + for handle in handles: + with self._lock: + tracked = self._works.get(id(handle)) is handle + if not tracked: + continue + with handle._state_lock: + phase = handle.phase + if phase == "recovery_pending": + self._retry_interrupted_post_recovery( + handle, + schedule=False, + ) + with handle._state_lock: + phase = handle.phase + if phase in ("posting", "recovery_pending"): + pending = True + continue + if phase == "prepared": + with handle._state_lock: + if handle.phase == "prepared": + handle.terminal_state = "ERR" + handle.phase = "terminal" + if self._work_state(handle, release=False) not in ( + "DONE", + "ERR", + ): + pending = True + continue + try: + self._release_work(handle) + except Exception as exc: + logger.warning( + "PyTorch failed-batch terminal cleanup failed: %s", + exc, + ) + pending = True + if not pending: + with self._lock: + self._prune_retained_handle_batches_unlocked() + return + time.sleep(0) + + def release_xfer_handle(self, handle: _WorkHandle) -> None: + with self._admitted_call(): + with self._lock: + tracked = self._works.get(id(handle)) is handle + if not tracked: + return + with handle._state_lock: + phase = handle.phase + if phase == "prepared": + self._release_work(handle) + return + if phase == "posting": + raise RuntimeError("Cannot release a posting PyTorch transfer") + if phase == "recovery_pending": + self._retry_interrupted_post_recovery(handle, schedule=True) + with handle._state_lock: + recovery_pending = handle.phase == "recovery_pending" + if recovery_pending: + raise RuntimeError( + "Cannot release a PyTorch transfer while Work recovery is pending" + ) + if self._work_state(handle, release=False) not in ("DONE", "ERR"): + with handle._state_lock: + work = handle.work + assert work is not None + work.cancel() + self._progress() + if self._work_state(handle, release=False) not in ("DONE", "ERR"): + raise RuntimeError("Cannot release an active PyTorch transfer") + self._release_work(handle) + + def send_notif(self, peer_name: str, payload: bytes) -> _WorkHandle | None: + with self._admitted_call() as endpoint: + with self._lock: + peer = self._require_peer(peer_name) + send_notification_sync = self._send_notification_sync + if type(payload) is not bytes: + payload = bytes(payload) + if send_notification_sync is not None: + send_notification_sync(peer, payload) + return None + handle = _WorkHandle(phase="posting") + try: + # Tracking and provider entry are one recovery transaction. + self._track_work(handle) + work = endpoint.send_notification( + peer, + payload, + recovery_token=handle, + ) + with handle._state_lock: + handle.work = work + handle.phase = "posted" + with self._lock: + self._detached_works.append(handle) + return handle + except BaseException as error: + self._handle_interrupted_submission(handle, error) + with self._lock: + tracked = self._works.get(id(handle)) is handle + if tracked: + self._ensure_detached(handle) + raise + + def _notification_alias( + self, + source_endpoint_id: str | None, + source_incarnation: str | None, + source_aliases: Mapping[object, str] | None = None, + ) -> str: + if source_endpoint_id is None: + if source_incarnation is not None: + raise RuntimeError("notification source identity is incomplete") + # One-way WRITE topologies need not import the sender. Notification + # payloads remain self-routing in SGLang's current protocol. + return _UNKNOWN_NOTIFICATION_SOURCE + if source_incarnation is None: + raise RuntimeError( + "notification source is missing its endpoint incarnation" + ) + source_key = (str(source_endpoint_id), str(source_incarnation)) + if source_aliases is None: + with self._lock: + source_aliases = self._source_aliases + try: + return source_aliases[source_key] + except KeyError as exc: + raise RuntimeError( + "notification arrived from an unknown or stale endpoint " + f"incarnation {source_key!r}" + ) from exc + + def get_new_notifs(self) -> dict[str, list[bytes]]: + with self._admitted_call() as endpoint: + poll_batches = getattr(endpoint, "poll_notification_batches", None) + if callable(poll_batches): + notifications = poll_batches() + grouped_receive = True + else: + notifications = endpoint.poll_notifications() + grouped_receive = False + # Either polling operation already drives NIXL notification progress. + # Reap detached notification Works without issuing a second + # provider progress call or scanning caller-owned transfers. + self._reap_detached_works() + # Peer mutation publishes a replacement alias map under ``_lock``. + # Snapshot only after native polling, then route the whole batch + # without a lock per notification. + with self._lock: + source_aliases = self._source_aliases + grouped: dict[str, list[bytes]] = defaultdict(list) + if grouped_receive: + for batch in notifications: + alias = self._notification_alias( + batch.source_endpoint_id, + batch.source_incarnation, + source_aliases, + ) + grouped[alias].extend(batch.payloads) + return dict(grouped) + if isinstance(notifications, Mapping): + for source, payloads in notifications.items(): + if source is None: + alias = _UNKNOWN_NOTIFICATION_SOURCE + else: + source = str(source) + alias = source_aliases.get(source, source) + grouped[alias].extend(bytes(payload) for payload in payloads) + return dict(grouped) + for notification in notifications: + alias = self._notification_alias( + notification.source_endpoint_id, + getattr(notification, "source_incarnation", None), + source_aliases, + ) + grouped[alias].append(bytes(notification.payload)) + return dict(grouped) + + def close(self) -> None: + endpoint = self._begin_close() + if endpoint is None: + return + try: + # Endpoint.closed is the authoritative commit receipt. Inspect it + # before querying children because a previous close may have + # committed even though its Python return was lost. + if self._endpoint_is_closed(endpoint): + self._finalize_closed_endpoint_unlocked(endpoint) + return + self._close_open_endpoint_exclusive(endpoint) + finally: + self._end_close() + + def _close_open_endpoint_exclusive(self, endpoint: object) -> None: + """Close one live endpoint while the lifecycle gate is exclusive.""" + + try: + self._retry_endpoint_cleanup_unlocked() + self._reconcile_construction_state_unlocked(force=True) + for handle in tuple(self._works.values()): + if handle.phase == "recovery_pending": + self._retry_interrupted_post_recovery(handle, schedule=False) + for handle in tuple(self._works.values()): + if handle.phase == "prepared": + self._release_work(handle) + for handle in tuple(self._works.values()): + if handle.phase == "posted" and self._work_state( + handle, release=False + ) not in ("DONE", "ERR"): + assert handle.work is not None + handle.work.cancel() + self._progress() + active = [ + handle + for handle in self._works.values() + if handle.phase in ("posting", "recovery_pending") + or ( + handle.phase == "posted" + and self._work_state(handle, release=False) not in ("DONE", "ERR") + ) + ] + if active: + raise RuntimeError( + f"Cannot close endpoint with {len(active)} active transfer(s)" + ) + for handle in tuple(self._works.values()): + self._release_work(handle) + self._detached_works.clear() + self._retained_work_ids.clear() + self._active_handle_batches.clear() + self._retained_handle_batches.clear() + self._retiring_peers.update(self._peers) + self._release_cached_plans(lambda cached: True) + for peer in reversed(tuple(self._peers.values())): + peer.close() + self._peers.clear() + self._peer_wire_metadata.clear() + self._retiring_peers.clear() + for record in reversed(self._registrations): + record.registration.close() + self._registrations.clear() + self._registration_groups.clear() + endpoint.close() + except BaseException: + if self._endpoint_is_closed(endpoint): + self._finalize_closed_endpoint_unlocked(endpoint) + raise + self._finalize_closed_endpoint_unlocked(endpoint) + + @contextmanager + def _admitted_call(self): + """Admit native work without holding a Python lock across the call.""" + + depth = getattr(self._admission_local, "depth", 0) + if depth: + self._admission_local.depth = depth + 1 + try: + yield self._require_endpoint() + finally: + self._admission_local.depth = depth + return + thread_id = threading.get_ident() + with self._lifecycle_cv: + if self._closing and self._close_owner != thread_id: + raise RuntimeError("PyTorch transfer endpoint close is in progress") + endpoint = self._require_endpoint() + self._active_calls += 1 + self._admission_local.depth = 1 + try: + yield endpoint + finally: + with self._lifecycle_cv: + self._admission_local.depth = 0 + self._active_calls -= 1 + if self._active_calls == 0: + self._lifecycle_cv.notify_all() + + def _begin_close(self) -> object | None: + thread_id = threading.get_ident() + with self._lifecycle_cv: + while self._closing and self._close_owner != thread_id: + self._lifecycle_cv.wait() + if self._close_owner == thread_id: + raise RuntimeError("recursive PyTorch transfer endpoint close") + if self._endpoint is None: + return None + self._closing = True + self._close_owner = thread_id + while self._active_calls: + self._lifecycle_cv.wait() + return self._endpoint + + def _end_close(self) -> None: + with self._lifecycle_cv: + self._closing = False + self._close_owner = None + self._lifecycle_cv.notify_all() + + @staticmethod + def _endpoint_is_closed(endpoint: object) -> bool: + try: + return getattr(endpoint, "closed", False) is True + except BaseException: + return False + + def _finalize_closed_endpoint_unlocked(self, endpoint: object) -> None: + """Adopt authoritative endpoint close without touching Core children.""" + + if self._endpoint is not endpoint: + return + for handle in self._works.values(): + if handle.terminal_state is None: + handle.terminal_state = "ERR" + handle.phase = "released" + handle.local_indices = None + handle.remote_indices = None + handle.notification = None + handle.plan = None + self._works.clear() + self._detached_works.clear() + self._retained_work_ids.clear() + self._active_handle_batches.clear() + self._retained_handle_batches.clear() + self._plan_cache.clear() + self._closing_plans.clear() + self._retiring_peers.clear() + self._peers.clear() + self._peer_wire_metadata.clear() + self._source_aliases.clear() + self._registrations.clear() + self._registration_groups.clear() + self._send_notification_sync = None + self._endpoint = None + self._construction_dirty = False + self._sweep_wakeup.set() + + @staticmethod + def _supports_capability(endpoint: object, name: str) -> bool: + capabilities = getattr(endpoint, "capabilities", None) + if callable(capabilities): + capabilities = capabilities() + if isinstance(capabilities, Mapping): + return bool(capabilities.get(name, False)) + return bool(getattr(capabilities, name, False)) + + @staticmethod + def _normalize_descriptors( + descriptors: Sequence[Sequence[object]], + ) -> tuple[_RawDescriptor, ...]: + normalized = [] + for descriptor in descriptors: + if len(descriptor) < 3: + raise ValueError("Transfer descriptors require address, size, device") + address, nbytes, device_id = map(int, descriptor[:3]) + if address <= 0 or nbytes <= 0 or device_id < 0: + raise ValueError( + "Transfer descriptor address/size must be positive and device " + "must be non-negative" + ) + stride: int | None = None + count: int | None = None + if len(descriptor) >= 5: + stride, count = map(int, descriptor[3:5]) + if stride < nbytes or count <= 0: + raise ValueError( + "Strided descriptors require stride >= size and count > 0" + ) + if count == 1: + stride = count = None + elif len(descriptor) == 4 and not isinstance(descriptor[3], str): + raise ValueError("Strided descriptors require both stride and count") + normalized.append(_RawDescriptor(address, nbytes, device_id, stride, count)) + return tuple(normalized) + + @staticmethod + def _device(memory_type: str, device_id: int) -> str: + if memory_type == "DRAM": + return "cpu" + if memory_type == "VRAM": + return f"cuda:{device_id}" + raise NotImplementedError( + f"PyTorch transfer adapter does not support {memory_type!r} memory" + ) + + def _resolve_local( + self, descriptors: Sequence[_RawDescriptor], memory_type: str + ) -> tuple[object, ...]: + regions = [] + for descriptor in descriptors: + end = descriptor.address + descriptor.extent + for record in self._registrations: + if ( + record.memory_type == memory_type + and record.device_id == descriptor.device_id + and descriptor.address >= record.address + and end <= record.address + record.nbytes + ): + kwargs = ( + {} + if descriptor.count is None + else { + "stride": descriptor.stride, + "count": descriptor.count, + } + ) + regions.append( + record.registration.region( + descriptor.address - record.address, + descriptor.nbytes, + **kwargs, + ) + ) + break + else: + raise ValueError( + f"Local raw span 0x{descriptor.address:x}+{descriptor.nbytes} " + "is outside registered memory" + ) + return tuple(regions) + + def _resolve_remote( + self, + peer_name: str, + descriptors: Sequence[_RawDescriptor], + memory_type: str, + ) -> tuple[object, ...]: + peer = self._require_peer(peer_name) + return tuple( + peer.unsafe_region( + descriptor.address, + descriptor.nbytes, + device=self._device(memory_type, descriptor.device_id), + memory_type=memory_type, + **( + {} + if descriptor.count is None + else { + "stride": descriptor.stride, + "count": descriptor.count, + } + ), + ) + for descriptor in descriptors + ) + + def _orient_prepared( + self, + operation: str, + source: _PreparedList, + source_indices: Sequence[int], + destination: _PreparedList, + destination_indices: Sequence[int], + ) -> tuple[_PreparedList, _PreparedList, Sequence[int], Sequence[int]]: + if operation == "WRITE": + local, remote = source, destination + local_indices, remote_indices = source_indices, destination_indices + else: + local, remote = destination, source + local_indices, remote_indices = destination_indices, source_indices + if local.remote or not remote.remote: + raise ValueError( + f"Invalid {operation} prepared-list local/remote orientation" + ) + return local, remote, local_indices, remote_indices + + def _post(self, handle: _WorkHandle, endpoint: object) -> None: + if handle.operation is None or handle.plan is None: + raise RuntimeError("PyTorch transfer handle has no prepared operation") + with self._lock: + if handle.peer_name in self._retiring_peers: + raise RuntimeError(f"Peer {handle.peer_name!r} teardown is in progress") + if handle.peer_name is not None: + self._require_peer(handle.peer_name) + if handle.cache_key is not None and handle.cache_key in self._closing_plans: + raise RuntimeError("Prepared transfer plan teardown is in progress") + with handle._state_lock: + handle._native_submit_entered = True + first_state = None + if handle.local_indices is not None or handle.remote_indices is not None: + if handle.local_indices is None or handle.remote_indices is None: + raise RuntimeError("PyTorch indexed transfer has incomplete indices") + kwargs = { + "local_indices": handle.local_indices, + "remote_indices": handle.remote_indices, + "max_polls": 0, + "timeout_ns": None, + "recovery_token": handle, + } + if handle.notification is None: + work, first_state = handle.plan.submit_indices_and_poll( + self._core_transfer_operations[handle.operation], + **kwargs, + ) + else: + work, first_state = handle.plan.submit_indices_and_poll( + self._core_transfer_operations[handle.operation], + **kwargs, + notification=handle.notification, + ) + else: + method = getattr(endpoint, handle.operation.lower()) + if handle.notification is None: + work = method( + handle.plan, + recovery_token=handle, + ) + else: + work = method( + handle.plan, + notification=handle.notification, + recovery_token=handle, + ) + with handle._state_lock: + handle.work = work + handle.first_state = first_state + # Core Work now owns its authenticated private snapshot. Drop the + # adapter's temporary immutable inputs on the successful boundary. + handle.local_indices = None + handle.remote_indices = None + handle.phase = "posted" + + def _recover_interrupted_post_work( + self, + handle: _WorkHandle, + submission_error: BaseException | None = None, + ) -> Literal["absent", "recovered", "pending"]: + """Recover one Core Work by an exact, caller-owned identity token.""" + + # Claim recovery with a state-only critical section. The Core ownership + # snapshot below must not run under either an agent or handle lock, and + # a second recovery caller must not race adoption of the same Work. + with handle._state_lock: + if handle.work is not None: + handle.phase = "posted" + return "recovered" + if handle.phase not in ("posting", "recovery_pending"): + return "absent" + if handle._recovery_in_progress: + return "pending" + handle.phase = "recovery_pending" + handle._recovery_in_progress = True + match = None + matches = () + try: + # This snapshot stays entirely off the successful path. Distractor + # Works cannot be adopted: token comparison is strictly ``is``. + matches = tuple( + work + for work in self._require_endpoint().open_works + if work.recovery_token is handle + ) + except BaseException as recovery_error: + self._note_recovery_issue(submission_error, recovery_error) + recovery = "pending" + else: + if not matches: + recovery = "absent" + elif len(matches) == 1: + match = matches[0] + recovery = "recovered" + else: + recovery = "pending" + if len(matches) > 1: + self._note_recovery_issue( + submission_error, + RuntimeError( + "multiple Core Works carry one SGLang recovery-token identity" + ), + ) + with handle._state_lock: + try: + if recovery == "recovered": + handle.work = match + handle.phase = "posted" + elif recovery == "absent": + # Absence is authoritative only for the recovery owner. Mark + # terminal before releasing the claim so no retry can race + # between the empty snapshot and terminal publication. + handle.terminal_state = "ERR" + handle.phase = "terminal" + finally: + handle._recovery_in_progress = False + return recovery + + def _handle_interrupted_submission( + self, + handle: _WorkHandle, + submission_error: BaseException, + ) -> None: + with handle._state_lock: + handle.submission_interrupted = ( + handle.phase + in ( + "posting", + "recovery_pending", + "posted", + ) + or handle.work is not None + ) + recovery = self._recover_interrupted_post_work(handle, submission_error) + if recovery == "absent": + self._close_failed_submission( + handle, + submission_error, + schedule=True, + ) + return + if recovery == "recovered": + self._quarantine_recovered_submission(handle, schedule=True) + return + self._retain_recovery_pending(handle, schedule=True) + + def _retry_interrupted_post_recovery( + self, + handle: _WorkHandle, + *, + schedule: bool, + ) -> None: + with handle._state_lock: + if handle.phase != "recovery_pending": + return + recovery = self._recover_interrupted_post_work(handle) + if recovery == "pending": + self._retain_recovery_pending(handle, schedule=schedule) + return + if recovery == "absent": + self._close_failed_submission(handle, schedule=schedule) + return + self._quarantine_recovered_submission(handle, schedule=schedule) + + def _quarantine_recovered_submission( + self, + handle: _WorkHandle, + *, + schedule: bool, + ) -> None: + with handle._state_lock: + work = handle.work + assert work is not None + try: + work.cancel() + except Exception as exc: + logger.warning("PyTorch interrupted Work cancellation failed: %s", exc) + if self._work_state(handle, release=False) in ("DONE", "ERR"): + self._close_failed_submission(handle, schedule=schedule) + return + self._retain_recovery_pending(handle, schedule=schedule) + + def _close_failed_submission( + self, + handle: _WorkHandle, + submission_error: BaseException | None = None, + *, + schedule: bool, + ) -> None: + try: + self._release_work(handle) + except BaseException as cleanup_error: + self._note_recovery_issue(submission_error, cleanup_error) + logger.warning("PyTorch failed-post cleanup failed: %s", cleanup_error) + self._retain_recovery_pending(handle, schedule=schedule) + if not isinstance(cleanup_error, Exception): + raise + + def _retain_recovery_pending( + self, + handle: _WorkHandle, + *, + schedule: bool, + ) -> None: + retained = False + with self._lock: + if self._works.get(id(handle)) is handle: + self._retained_work_ids.add(id(handle)) + retained = True + if schedule and retained: + self._start_retained_work_sweeper() + + def _require_no_unsettled_interrupted_submission(self, action: str) -> None: + """Fail closed using state only; callers may already own ``_lock``. + + Recovery and status observation are driven by admitted transfer or + sweeper calls. A teardown path must never hide either Core operation + beneath the agent RLock merely to make cleanup opportunistically pass. + """ + + unsettled = [] + for handle in tuple(self._works.values()): + with handle._state_lock: + phase = handle.phase + submission_interrupted = handle.submission_interrupted + if phase in ("posting", "recovery_pending"): + unsettled.append(handle) + continue + if not submission_interrupted or phase not in ( + "posted", + "terminal", + ): + continue + unsettled.append(handle) + if unsettled: + raise RuntimeError( + f"Cannot {action} with {len(unsettled)} interrupted PyTorch " + "submission(s) unsettled" + ) + + def _ensure_detached(self, handle: _WorkHandle) -> None: + with self._lock: + if self._works.get(id(handle)) is not handle: + return + if not any(candidate is handle for candidate in self._detached_works): + self._detached_works.append(handle) + + @staticmethod + def _note_recovery_issue( + submission_error: BaseException | None, + recovery_error: BaseException, + ) -> None: + if submission_error is None: + return + add_note = getattr(submission_error, "add_note", None) + if callable(add_note): + try: + add_note( + "PyTorch interrupted-post Work recovery reported " + f"{type(recovery_error).__qualname__}" + ) + except BaseException: # noqa: S110 + pass + + def _track_work(self, handle: _WorkHandle) -> _WorkHandle: + batch_handles = getattr(self._batch_handles, "value", None) + try: + with self._lock: + if batch_handles is not None and handle.operation is not None: + # This publication precedes provider entry. If a later helper + # frame is interrupted, the room-level barrier still owns it. + batch_handles.append(handle) + # Publish to the graph only after a worker batch can reach it. + self._works[id(handle)] = handle + return handle + except BaseException: + with self._lock: + self._works.pop(id(handle), None) + if batch_handles is not None: + batch_handles[:] = [ + candidate + for candidate in batch_handles + if candidate is not handle + ] + raise + + def _progress(self) -> None: + self._advance_provider() + self._reap_works() + + def _advance_provider(self) -> None: + endpoint = self._require_endpoint() + if not self._background_progress: + try: + endpoint.progress() + except Exception as exc: + # Global progress is an observation boundary too: failure does + # not prove that any request is terminal. Subsequent per-Work + # probes retain every unknown request and recover independently. + logger.warning( + "PyTorch transfer progress is temporarily unknown: %s", exc + ) + + def _reap_works(self) -> None: + with self._lock: + handles = tuple(self._works.values()) + self._reap_handles(handles) + + def _reap_detached_works(self) -> None: + with self._lock: + handles = tuple(self._detached_works) + if not handles: + return + self._reap_handles(handles) + + def _reap_retained_works(self) -> None: + retained = [] + seen = set() + with self._lock: + for handle_id in self._retained_work_ids: + handle = self._works.get(handle_id) + if handle is not None: + retained.append(handle) + seen.add(handle_id) + for handles in self._retained_handle_batches.values(): + for handle in handles: + handle_id = id(handle) + if handle_id in seen or self._works.get(handle_id) is not handle: + continue + retained.append(handle) + seen.add(handle_id) + for handle in retained: + with handle._state_lock: + if handle.phase == "prepared": + handle.terminal_state = "ERR" + handle.phase = "terminal" + self._reap_handles(tuple(retained)) + with self._lock: + self._prune_retained_handle_batches_unlocked() + + def _reap_handles(self, handles: Sequence[_WorkHandle]) -> None: + for handle in handles: + with self._lock: + tracked = self._works.get(id(handle)) is handle + if not tracked: + continue + with handle._state_lock: + phase = handle.phase + if phase == "recovery_pending": + self._retry_interrupted_post_recovery(handle, schedule=False) + with handle._state_lock: + phase = handle.phase + if phase in ("prepared", "posting", "recovery_pending"): + continue + if self._work_state(handle, release=False) not in ("DONE", "ERR"): + continue + try: + self._release_work(handle) + except Exception as exc: + # The Work remains tracked, and the next progress call retries + # cleanup without losing its terminal state. + logger.warning("PyTorch terminal Work cleanup failed: %s", exc) + with self._lock: + self._detached_works = [ + handle + for handle in self._detached_works + if self._works.get(id(handle)) is handle + ] + self._retained_work_ids.intersection_update(self._works) + + def _start_retained_work_sweeper(self) -> None: + with self._lock: + if self._sweep_thread is not None and self._sweep_thread.is_alive(): + self._sweep_wakeup.set() + return + self._sweep_wakeup.clear() + self._sweep_thread = threading.Thread( + target=self._sweep_retained_works, + name=f"{self.name}-retained-work-sweeper", + daemon=True, + ) + self._sweep_thread.start() + + def _sweep_retained_works(self) -> None: + """Periodically reap abandoned siblings without freeing active Works.""" + + while True: + self._sweep_wakeup.wait(_RETAINED_WORK_SWEEP_INTERVAL_S) + self._sweep_wakeup.clear() + with self._lock: + if self._endpoint is None: + self._sweep_thread = None + return + self._retained_work_ids.intersection_update(self._works) + self._prune_retained_handle_batches_unlocked() + if not self._retained_work_ids and not self._retained_handle_batches: + self._sweep_thread = None + return + try: + # Abandoned siblings need autonomous provider progress, but no + # adapter lock spans provider progress or Work-state refresh. + with self._admitted_call(): + self._advance_provider() + self._reap_retained_works() + self._reap_detached_works() + except RuntimeError: + with self._lock: + if self._closing or self._endpoint is None: + self._sweep_thread = None + return + raise + + def _prune_retained_handle_batches_unlocked(self) -> None: + for batch_id, handles in tuple(self._retained_handle_batches.items()): + if not any(self._works.get(id(handle)) is handle for handle in handles): + self._retained_handle_batches.pop(batch_id, None) + + def _work_state(self, handle: _WorkHandle, *, release: bool) -> str: + while True: + with handle._state_lock: + terminal_state = handle.terminal_state + phase = handle.phase + if terminal_state is not None: + break + if phase == "prepared": + raise RuntimeError("PyTorch transfer has not been posted") + if phase == "posting": + raise RuntimeError( + "PyTorch transfer submission is still in progress" + ) + if phase == "recovery_pending": + recover = True + else: + recover = False + if handle.work is None: + raise RuntimeError("Posted PyTorch transfer has no Work") + if handle._status_in_progress: + return "PROC" + handle._status_in_progress = True + work = handle.work + state = handle.first_state + handle.first_state = None + if not recover: + break + self._retry_interrupted_post_recovery(handle, schedule=True) + with handle._state_lock: + recovery_pending = handle.phase == "recovery_pending" + if recovery_pending: + return "PROC" + + if terminal_state is not None: + if release: + self._release_work(handle) + return terminal_state + try: + # Work.state is the sole authoritative observation. In current + # Core it refreshes once and already incorporates provider errors; + # reading test(), error, and state separately would refresh the + # same Work up to three times. + if state is None: + state = work.state + outcomes = self._work_state_outcomes + if outcomes is not None: + try: + terminal = outcomes.get(state, _UNKNOWN_WORK_OUTCOME) + except TypeError: + terminal = _UNKNOWN_WORK_OUTCOME + if terminal is _UNKNOWN_WORK_OUTCOME: + raise ValueError(f"Unknown PyTorch WorkState: {state!r}") + if terminal is None: + return "PROC" + else: + state_name = self._state_name(state) + if state_name in _PENDING_STATES: + return "PROC" + if state_name in _DONE_STATES: + terminal = "DONE" + elif state_name in _FAILED_STATES: + terminal = "ERR" + else: + raise ValueError(f"Unknown PyTorch WorkState: {state!r}") + except Exception as exc: + # A status observation failure does not prove the transport is + # terminal. Returning PROC preserves all memory/plan lifetimes and + # permits a later poll to recover or report a real terminal error. + logger.warning("PyTorch transfer status is temporarily unknown: %s", exc) + return "PROC" + finally: + with handle._state_lock: + handle._status_in_progress = False + with handle._state_lock: + if handle.submission_interrupted: + terminal = "ERR" + if handle.terminal_state is None: + handle.terminal_state = terminal + handle.phase = "terminal" + else: + terminal = handle.terminal_state + if release: + self._release_work(handle) + return terminal + + @staticmethod + def _complete_work_state_outcomes( + api: Any, + ) -> dict[object, str | None] | None: + work_state = getattr(api, "WorkState", None) + try: + outcomes = { + work_state.PENDING: None, + work_state.RUNNING: None, + work_state.COMPLETED: "DONE", + work_state.FAILED: "ERR", + work_state.CANCELLED: "ERR", + } + except (AttributeError, TypeError): + return None + return outcomes if len(outcomes) == 5 else None + + def _release_work(self, handle: _WorkHandle) -> None: + with self._lock: + tracked = self._works.get(id(handle)) is handle + if not tracked: + return + with handle._state_lock: + if handle.phase == "released" or handle._release_in_progress: + return + if handle.phase in ("posting", "recovery_pending"): + raise RuntimeError( + "Cannot release a Work with unresolved submission ownership" + ) + handle._release_in_progress = True + work = handle.work + owned_plan = handle.owned_plan + try: + if work is not None: + work.close() + if owned_plan is not None: + owned_plan.close() + except BaseException: + with handle._state_lock: + handle._release_in_progress = False + raise + with self._lock: + if self._works.get(id(handle)) is handle: + self._works.pop(id(handle), None) + if handle.operation is None: + self._detached_works = [ + candidate + for candidate in self._detached_works + if candidate is not handle + ] + with handle._state_lock: + handle.phase = "released" + handle.local_indices = None + handle.remote_indices = None + handle.notification = None + handle.plan = None + handle._release_in_progress = False + + def _release_cached_plans(self, predicate: Callable[[_CachedPlan], bool]) -> None: + # Move every matching plan out of the submit-visible cache before the + # first close attempt. A partially committed provider close remains + # discoverable here and may be retried, but can never be submitted. + for key, cached in tuple(self._plan_cache.items()): + if not predicate(cached): + continue + del self._plan_cache[key] + self._closing_plans[key] = cached + for key, cached in reversed(tuple(self._closing_plans.items())): + if not predicate(cached): + continue + if any(handle.cache_key == key for handle in self._works.values()): + raise RuntimeError( + "Prepared transfer plan is still referenced by a transfer handle" + ) + cached.plan.close() + del self._closing_plans[key] + + @staticmethod + def _contains_identity(values: Sequence[object], candidate: object) -> bool: + return any(value is candidate for value in values) + + def _reconcile_construction_state_unlocked(self, *, force: bool = False) -> None: + """Close exact Core children absent from this adapter's owner graph.""" + + if not force and not self._construction_dirty: + return + endpoint = self._require_endpoint() + entries = tuple(self._plan_cache.values()) + tuple(self._closing_plans.values()) + owned_plans = tuple(entry.plan for entry in entries) + tuple( + handle.owned_plan + for handle in self._works.values() + if handle.owned_plan is not None + ) + owned_peers = tuple(self._peers.values()) + owned_registrations = tuple( + record.registration for record in self._registrations + ) + # SGLang does not currently construct bindings or request slots. They + # are still scanned first so any future interrupted child cannot pin a + # lost plan while being silently mistaken for an owned object. + for attribute, owned in ( + ("open_transfer_slots", ()), + ("open_bound_transfers", ()), + ("open_plans", owned_plans), + ("open_peers", owned_peers), + ("open_registrations", owned_registrations), + ): + for child in tuple(getattr(endpoint, attribute)): + if not self._contains_identity(owned, child): + child.close() + self._construction_dirty = False + + def _reconcile_after_construction_error_unlocked( + self, error: BaseException, kind: str + ) -> None: + try: + self._reconcile_construction_state_unlocked(force=True) + except BaseException as cleanup_error: + self._add_exception_note( + error, + f"{kind} construction cleanup also raised: " + f"{type(cleanup_error).__qualname__}", + ) + + def _matches_recovery_endpoint( + self, + candidate: object, + *, + endpoint_id: str, + progress_mode: object, + thread_mode: object, + ) -> bool: + return ( + type(candidate) is self._api.Endpoint + and getattr(candidate, "endpoint_id", None) == endpoint_id + and getattr(candidate, "name", None) == self.name + and getattr(candidate, "backend", None) == self._endpoint_backend + and getattr(candidate, "progress_mode", None) is progress_mode + and getattr(candidate, "thread_mode", None) is thread_mode + ) + + @staticmethod + def _close_recovery_endpoint(endpoint: object) -> None: + if not getattr(endpoint, "closed", False): + endpoint.close() + if not getattr(endpoint, "closed", False): + raise RuntimeError("Endpoint cleanup returned before close committed") + + def _retry_endpoint_cleanup_unlocked(self) -> None: + endpoint = self._endpoint_recovery + if endpoint is not None: + self._close_recovery_endpoint(endpoint) + self._endpoint_recovery = None + + def _retry_endpoint_recovery_unlocked(self) -> None: + """Compatibility alias for focused lifecycle tests.""" + + self._retry_endpoint_cleanup_unlocked() + + @staticmethod + def _add_exception_note(error: BaseException, message: str) -> None: + add_note = getattr(error, "add_note", None) + if callable(add_note): + try: + add_note(message) + except BaseException: # noqa: S110 + pass + + @staticmethod + def _state_name(state: object) -> str | None: + state = getattr(state, "name", state) + return state.upper() if isinstance(state, str) else None + + @staticmethod + def _normalize_operation(operation: str) -> Literal["READ", "WRITE"]: + operation = operation.upper() + if operation not in ("READ", "WRITE"): + raise ValueError(f"Unsupported transfer operation: {operation!r}") + return operation # type: ignore[return-value] + + @staticmethod + def _peer_source_ids(peer: object) -> tuple[str, ...]: + identities = [] + for attribute in ("endpoint_id", "name", "source_endpoint_id"): + value = getattr(peer, attribute, None) + if value is not None: + identity = str(value) + if identity not in identities: + identities.append(identity) + if not identities: + raise TypeError("Imported peer does not expose an endpoint identity") + return tuple(identities) + + def _require_endpoint(self) -> Any: + if self._closing and self._close_owner != threading.get_ident(): + raise RuntimeError("PyTorch transfer endpoint close is in progress") + if self._endpoint is None: + raise RuntimeError("PyTorch transfer endpoint is not initialized") + return self._endpoint + + def _require_peer(self, peer_name: str) -> Any: + self._require_endpoint() + try: + return self._peers[peer_name] + except KeyError as exc: + raise KeyError(f"Unknown PyTorch transfer peer {peer_name!r}") from exc diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index d4163ea8835c..a062e1c4b354 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -1191,6 +1191,14 @@ def maybe_send_cached_prefix_chunk(self: Scheduler, req: Req) -> None: if not envs.SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX.get(): return + # Early-send is safe only when the selected data plane can defer source + # access until the transfer worker consumes the recorded CUDA event. + # Keep the capability gate explicit for alternate sender implementations. + if self.enable_overlap and not getattr( + req.disagg_kv_sender, "supports_overlap_early_send", True + ): + return + # Staging sends into positional grid slots, so the early-send boundary # must stay stable across the request's batches: snapshot the at-rest # prefix on the first batch. Non-staging reads the live prefix. diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index a5e613f08c54..272fe38f01a4 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -669,16 +669,18 @@ class Envs: SGLANG_DISAGGREGATION_ENGINE_INIT_TIMEOUT = EnvInt(60) SGLANG_DISAGGREGATION_NIXL_BACKEND = EnvStr("UCX") SGLANG_DISAGGREGATION_NIXL_BACKEND_PARAMS = EnvStr("{}") + SGLANG_DISAGGREGATION_NIXL_USE_TORCH_TRANSFER = EnvBool(False) SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX = EnvBool(True) SGLANG_DISAGGREGATION_ZMQ_MAX_SOCKETS = EnvInt(16384) SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER = EnvBool(False) SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK = EnvBool(False) SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS = EnvInt(0) SGLANG_DISAGGREGATION_BOOTSTRAP_ENTRY_CLEANUP_INTERVAL = EnvInt(120) - # Deferred decode-side KV release: on abort, hold an in-flight request's KV - # pages/slot until the prefill acks the transfer drained, or the timeout - # below fires. Off by default (no behavior/perf impact when disabled). - SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE = EnvBool(False) + # Compatibility knob retained for configuration parsing. Decode-side KV + # quarantine is now a mandatory correctness property; a false value is + # ignored by the transfer implementation. The timeout controls only the + # rate-limited operator diagnostic and ABORT retry, never page release. + SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE = EnvBool(True) SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE_TIMEOUT = EnvFloat(30.0) # =================================================================== diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 764205f08a53..ad105842f341 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1068,6 +1068,10 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True): # Cache namespace used to isolate otherwise-identical prefixes. cache_salt: Optional[str] = None + # Internal 128-bit room-lease identity, minted once before scheduler fanout. + # Keep this append-only: TokenizedGenerateReqInput is an array-like wire type. + transfer_generation: Optional[str] = None + def wrap_pickle_fields(self): self.time_stats = wrap_as_pickle(self.time_stats) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index d5e85b986d0f..e1173e665fe2 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -957,6 +957,7 @@ def __init__( disagg_mode: Optional[DisaggregationMode] = None, routed_dp_rank: Optional[int] = None, disagg_prefill_dp_rank: Optional[int] = None, + transfer_generation: Optional[str] = None, vocab_size: Optional[int] = None, priority: Optional[int] = None, metrics_collector: Optional[SchedulerMetricsCollector] = None, @@ -1269,6 +1270,9 @@ def __init__( self.bootstrap_host: str = bootstrap_host self.bootstrap_port: Optional[int] = bootstrap_port self.bootstrap_room: Optional[int] = bootstrap_room + # Shared by every scheduler rank for this decode-room lease. Unlike + # bootstrap_room and rid, this internal value is never user-recycled. + self.transfer_generation: Optional[str] = transfer_generation # Decode-local: the already-emitted boundary token to replay when a # retracted request is rebootstrapped. Set in pause_generation(retract) # and consumed in the decode transfer commit; never plumbed to prefill. diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 85b9c203c950..c6a67140fdb4 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2761,6 +2761,7 @@ def handle_generate_request( disagg_mode=self.disaggregation_mode, routed_dp_rank=recv_req.routed_dp_rank, disagg_prefill_dp_rank=recv_req.disagg_prefill_dp_rank, + transfer_generation=recv_req.transfer_generation, vocab_size=self.model_config.vocab_size, priority=recv_req.priority, metrics_collector=( @@ -5239,20 +5240,11 @@ def abort_request(self, recv_req: AbortReq): if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid): logger.debug(f"Abort transfer queue request. {decode_req.req.rid=}") receiver = decode_req.kv_receiver + # Publish exact-generation quarantine before the fallible + # ABORT send. A BaseException/lost return cannot race an ACK + # ahead of decode ownership. + receiver.arm_abort_intent() receiver.abort() - # Arm drain-ack accounting once the ABORT is sent, so acks - # arriving before this req is deferred (e.g. during the next - # forward step) are captured. A fresh set also drops stale acks - # from a prior request that reused this bootstrap_room. A - # redundant abort only re-wipes -- holds longer, never releases - # early -- so no transition guard is needed. - if ( - receiver.kv_mgr.enable_deferred_decode_kv_release - and receiver.abort_notified - ): - receiver.kv_mgr.register_deferred_abort_room( - decode_req.req.bootstrap_room - ) # Abort requests whose KV is already backed up for retraction. if self.disagg_decode_prealloc_queue.retracted_queue: diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index ea77294d2f88..1b3f9a90b88b 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -22,6 +22,7 @@ import logging import os import pickle +import secrets import signal import socket import sys @@ -1426,6 +1427,9 @@ def _create_tokenized_object( return_indexer_topk=obj.return_indexer_topk, routed_dp_rank=obj.routed_dp_rank, disagg_prefill_dp_rank=obj.disagg_prefill_dp_rank, + transfer_generation=( + secrets.token_hex(16) if bootstrap_room is not None else None + ), priority=obj.priority, extra_key=obj.extra_key, cache_salt=obj.cache_salt, diff --git a/rust/sglang-server/src/message/io_struct.rs b/rust/sglang-server/src/message/io_struct.rs index cd8c8fe34842..9cc7fd494a0e 100644 --- a/rust/sglang-server/src/message/io_struct.rs +++ b/rust/sglang-server/src/message/io_struct.rs @@ -43,9 +43,9 @@ wire_struct! { lora_id: (), custom_logit_processor: (), positional_embed_overrides: (), - /// PD-disaggregation block — the last fields emitted; everything after - /// `disagg_prefill_dp_rank` in Python has a msgspec default and is - /// omitted (short arrays decode with defaulted tails). + /// PD-disaggregation block. The internal generation is append-only at + /// the end of Python's array-like schema, so all intervening defaulted + /// fields must remain present here to preserve their wire positions. bootstrap_host: Option<&'a str>, bootstrap_port: Option, bootstrap_room: Option, @@ -53,6 +53,21 @@ wire_struct! { decode_tp_size: Option, routed_dp_rank: Option, disagg_prefill_dp_rank: Option, + routing_key: (), + require_reasoning: bool, + priority: (), + extra_key: (), + no_logs: bool, + return_bytes: bool, + return_entropy: bool, + need_wait_for_mm_inputs: (), + num_items_assigned: (), + encoder_urls: (), + multi_item_delimiter_indices: (), + time_stats: (), + cache_salt: (), + /// Internal room-lease identity shared by every scheduler rank. + transfer_generation: Option<&'a str>, } } @@ -114,6 +129,20 @@ impl<'a> From<&'a GenerateRequest> for TokenizedGenerateReqInput<'a> { decode_tp_size: req.decode_tp_size, routed_dp_rank: req.routed_dp_rank, disagg_prefill_dp_rank: req.disagg_prefill_dp_rank, + routing_key: (), + require_reasoning: false, + priority: (), + extra_key: (), + no_logs: false, + return_bytes: false, + return_entropy: false, + need_wait_for_mm_inputs: (), + num_items_assigned: (), + encoder_urls: (), + multi_item_delimiter_indices: (), + time_stats: (), + cache_salt: (), + transfer_generation: req.transfer_generation.as_deref(), } } } @@ -180,9 +209,9 @@ mod tests { let bytes = TokenizedGenerateReqInput::from(&req).encode().unwrap(); let val = rmpv::decode::read_value(&mut &bytes[..]).unwrap(); let arr = val.as_array().expect("array"); - // msgspec requires >= 14 (through `stream`); we emit 32 (through - // `disagg_prefill_dp_rank`). Trailing defaulted fields are omitted. - assert_eq!(arr.len(), 32, "header ends at disagg_prefill_dp_rank"); + // msgspec requires >= 14 (through `stream`); generation is an + // append-only field at 45, after every pre-existing optional field. + assert_eq!(arr.len(), 46, "header ends at transfer_generation"); assert_eq!(arr[0].as_str(), Some("TokenizedGenerateReqInput")); assert_eq!(arr[1].as_str(), Some("r1")); assert!(arr[5].is_nil(), "idx 5 must be input_embeds (nil)"); @@ -210,6 +239,7 @@ mod tests { Some(true), "return_hidden_states at idx 16" ); + assert!(arr[45].is_nil(), "generation is nil without a PD room"); } /// The PD block must land on Python's wire indices 25–31, with the filler @@ -227,6 +257,7 @@ mod tests { decode_tp_size: Some(2), routed_dp_rank: Some(3), disagg_prefill_dp_rank: Some(4), + transfer_generation: Some("11111111111111111111111111111111".into()), ..Default::default() }; let bytes = TokenizedGenerateReqInput::from(&req).encode().unwrap(); @@ -245,5 +276,16 @@ mod tests { assert_eq!(arr[29].as_i64(), Some(2), "decode_tp_size at 29"); assert_eq!(arr[30].as_i64(), Some(3), "routed_dp_rank at 30"); assert_eq!(arr[31].as_i64(), Some(4), "disagg_prefill_dp_rank at 31"); + for i in [32, 34, 35, 39, 40, 41, 42, 43, 44] { + assert!(arr[i].is_nil(), "defaulted optional field at idx {i}"); + } + for i in [33, 36, 37, 38] { + assert_eq!(arr[i].as_bool(), Some(false), "false default at idx {i}"); + } + assert_eq!( + arr[45].as_str(), + Some("11111111111111111111111111111111"), + "transfer_generation at append-only idx 45" + ); } } diff --git a/rust/sglang-server/src/message/request.rs b/rust/sglang-server/src/message/request.rs index 10df74b94108..1b0d22316d21 100644 --- a/rust/sglang-server/src/message/request.rs +++ b/rust/sglang-server/src/message/request.rs @@ -7,6 +7,7 @@ use std::sync::LazyLock; use bytes::Bytes; use itertools::izip; use serde::Deserialize; +use uuid::Uuid; use super::io_struct::{ControlRequest, TokenizedGenerateReqInput}; use super::response::ResponseSink; @@ -30,6 +31,20 @@ use crate::utils::{environ::env_i64, error::Error}; static MAX_BATCH_REQS_PER_HTTP_REQ: LazyLock = LazyLock::new(|| env_i64("SGLANG_MAX_BATCH_REQS_PER_HTTP_REQ", 4096)); +/// Mint a full-width, OS-seeded 128-bit lease identity without adding another +/// random-number dependency. UUID v4 supplies 122 random bits and overwrites +/// six RFC version/variant bits; fill those exact positions from independent +/// random bytes of a second UUID before rendering canonical lowercase hex. +fn new_transfer_generation() -> String { + let primary = Uuid::new_v4(); + let refill = Uuid::new_v4(); + let mut bytes = *primary.as_bytes(); + let refill_bytes = refill.as_bytes(); + bytes[6] = (bytes[6] & 0x0f) | (refill_bytes[0] & 0xf0); + bytes[8] = (bytes[8] & 0x3f) | (refill_bytes[1] & 0xc0); + format!("{:032x}", u128::from_be_bytes(bytes)) +} + fn batch_size_exceeds_limit(batch_size: usize, limit: i64) -> bool { limit >= 0 && batch_size as u128 > limit as u128 } @@ -414,6 +429,9 @@ impl GenerateBody { return_text_in_logprobs, bootstrap_host, bootstrap_port, + transfer_generation: bootstrap_room + .as_ref() + .map(|_| new_transfer_generation()), bootstrap_room, bootstrap_pair_key, decode_tp_size, @@ -649,6 +667,8 @@ pub struct GenerateRequest { /// so these are pure passthrough for the scheduler/LB protocol. pub routed_dp_rank: Option, pub disagg_prefill_dp_rank: Option, + /// Internal random room-lease identity, never accepted from the client. + pub transfer_generation: Option, /// Multimodal inputs, carried opaquely. Consumed by the Encoding stage, /// which ships them to the MM worker pool; never read by the tokenizer or /// serialized onto the scheduler header. Boxed so the common text-only @@ -1291,7 +1311,16 @@ mod tests { assert_eq!(p.bootstrap_port, Some(8998)); assert_eq!(p.bootstrap_room, Some(7 + i as i64)); assert_eq!(p.routed_dp_rank, Some(1)); + let generation = p.transfer_generation.as_deref().unwrap(); + assert_eq!(generation.len(), 32); + assert!( + generation + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)), + "generation must be canonical lowercase 128-bit hex" + ); } + assert_ne!(ps[0].transfer_generation, ps[1].transfer_generation); let (ps, _) = requests( r#"{"text": ["a", "b"], "bootstrap_host": ["h1", "h2"], @@ -1302,6 +1331,7 @@ mod tests { assert_eq!(ps[1].bootstrap_host.as_deref(), Some("h2")); assert_eq!(ps[0].bootstrap_room, Some(10)); assert_eq!(ps[1].bootstrap_room, Some(20)); + assert_ne!(ps[0].transfer_generation, ps[1].transfer_generation); let err = requests(r#"{"text": ["a", "b"], "bootstrap_room": [1, 2, 3]}"#).unwrap_err(); assert!(err.to_string().contains("bootstrap_room"), "{err}"); 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 0a8936f219d6..cbc0fe720090 100644 --- a/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py +++ b/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py @@ -3,18 +3,24 @@ When a decode request is aborted while its prefill->decode KV transfer may still be in flight, the decode side holds its KV pages / req-slot instead of freeing them immediately (which could let the still-in-flight write land on pages already -reused by another request). The pages are released once every prefill rank acks -that its transfer drained (CommonKVManager.is_abort_release_safe), or a timeout -fires. See DecodeTransferQueue.resolve_deferred_releases. +reused by another request). The pages are released only once every exact prefill +rank for the immutable room generation acks that its transfer drained. A timeout +is diagnostic and never authorizes reuse. """ +import os +import threading import unittest from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import MagicMock, patch from sglang.srt.disaggregation import decode as decode_mod -from sglang.srt.disaggregation.common.conn import CommonKVManager -from sglang.srt.disaggregation.decode import DecodeTransferQueue +from sglang.srt.disaggregation.common.conn import ( + CommonKVManager, + next_transfer_generation, +) +from sglang.srt.disaggregation.decode import DecodePreallocQueue, DecodeTransferQueue +from sglang.srt.disaggregation.utils import TransferBackend from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -26,91 +32,91 @@ def _make_manager(): touch (avoids the heavy real __init__).""" mgr = CommonKVManager.__new__(CommonKVManager) mgr._deferred_abort_ack_tracker = {} + mgr._abort_state_lock = threading.Lock() return mgr +GEN_A = "11" * 16 +GEN_B = "22" * 16 + + 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) - self.assertFalse(mgr.is_abort_release_safe(room, required_acks=2)) + mgr.register_deferred_abort_room(room, GEN_A, frozenset({4, 9})) + self.assertFalse(mgr.is_abort_release_safe(room, GEN_A)) - mgr.note_abort_ack(room, 0) - self.assertFalse(mgr.is_abort_release_safe(room, required_acks=2)) + self.assertTrue(mgr.note_abort_ack(room, GEN_A, 4)) + self.assertFalse(mgr.is_abort_release_safe(room, GEN_A)) - mgr.note_abort_ack(room, 1) - self.assertTrue(mgr.is_abort_release_safe(room, required_acks=2)) + self.assertTrue(mgr.note_abort_ack(room, GEN_A, 9)) + self.assertTrue(mgr.is_abort_release_safe(room, GEN_A)) 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 - # 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)) + mgr.register_deferred_abort_room(room, GEN_A, frozenset({0, 1})) + mgr.note_abort_ack(room, GEN_A, 0) + mgr.note_abort_ack(room, GEN_A, 0) + self.assertFalse(mgr.is_abort_release_safe(room, GEN_A)) - def test_single_rank_fast_path(self): + def test_invalid_and_unexpected_rank_identities_are_rejected(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)) + mgr.register_deferred_abort_room(room, GEN_A, frozenset({7})) + for bad_rank in (True, -1, 8, 2**63): + self.assertFalse(mgr.note_abort_ack(room, GEN_A, bad_rank)) + self.assertFalse(mgr.is_abort_release_safe(room, GEN_A)) + self.assertTrue(mgr.note_abort_ack(room, GEN_A, 7)) + self.assertTrue(mgr.is_abort_release_safe(room, GEN_A)) def test_clear_deferred_abort_state(self): mgr = _make_manager() room = 103 - mgr.register_deferred_abort_room(room) - mgr.note_abort_ack(room, 0) - 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)) + mgr.register_deferred_abort_room(room, GEN_A, frozenset({0})) + mgr.note_abort_ack(room, GEN_A, 0) + mgr.clear_deferred_abort_state(room, GEN_A) + self.assertNotIn((room, GEN_A), mgr._deferred_abort_ack_tracker) + self.assertFalse(mgr.is_abort_release_safe(room, GEN_A)) def test_ack_before_register_is_dropped(self): - # An ack for a room that isn't actively held must not be recorded (it - # would otherwise pollute a later request reusing the same room). mgr = _make_manager() room = 104 - mgr.note_abort_ack(room, 0) # 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). + self.assertFalse(mgr.note_abort_ack(room, GEN_A, 0)) + self.assertNotIn((room, GEN_A), mgr._deferred_abort_ack_tracker) + + def test_room_reuse_is_generation_isolated(self): 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): + expected = frozenset({0, 1}) + mgr.register_deferred_abort_room(room, GEN_A, expected) + mgr.note_abort_ack(room, GEN_A, 0) + mgr.clear_deferred_abort_state(room, GEN_A) + mgr.register_deferred_abort_room(room, GEN_B, expected) + + self.assertFalse(mgr.note_abort_ack(room, GEN_A, 1)) + mgr.note_abort_ack(room, GEN_B, 0) + self.assertFalse(mgr.is_abort_release_safe(room, GEN_B)) + mgr.note_abort_ack(room, GEN_B, 1) + self.assertTrue(mgr.is_abort_release_safe(room, GEN_B)) + + def test_rearming_is_idempotent_but_expected_set_is_immutable(self): mgr = _make_manager() room = 106 - mgr.register_deferred_abort_room(room) - mgr.note_abort_ack(room, 0) - mgr.note_abort_ack(room, 1) - 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) - self.assertFalse(mgr.is_abort_release_safe(room, required_acks=2)) + expected = frozenset({0, 1}) + mgr.register_deferred_abort_room(room, GEN_A, expected) + mgr.note_abort_ack(room, GEN_A, 0) + mgr.register_deferred_abort_room(room, GEN_A, expected) + self.assertFalse(mgr.is_abort_release_safe(room, GEN_A)) + with self.assertRaises(RuntimeError): + mgr.register_deferred_abort_room(room, GEN_A, frozenset({0})) + + def test_expected_rank_set_cannot_be_empty_or_malformed(self): + mgr = _make_manager() + for bad in (frozenset(), frozenset({-1}), frozenset({True})): + with self.assertRaises(ValueError): + mgr.register_deferred_abort_room(107, GEN_A, bad) class _FakeIdxAllocator: @@ -133,16 +139,30 @@ def _make_queue(timeout=30.0): return q -def _make_decode_req(room, idx, mgr, n_prefill_ranks=1): +def _make_decode_req(room, idx, mgr, ranks=(0,), generation=GEN_A): receiver = SimpleNamespace( kv_mgr=mgr, - # One entry per prefill rank the decode notified of the abort; its length - # is the required drain-ack count (see DecodeTransferQueue._defer_release). - bootstrap_infos=[{"rank": r} for r in range(n_prefill_ranks)], - clear=lambda: None, + bootstrap_infos=[{"prefill_unique_rank": rank} for rank in ranks], + transfer_generation=generation, + abort_notified=False, + clear=MagicMock(), + retry_abort_notification=MagicMock(), + failure_exception=MagicMock(side_effect=RuntimeError("transfer failed")), ) + + def arm_abort_intent(): + if receiver.abort_notified: + return + mgr.register_deferred_abort_room(room, generation, frozenset(ranks)) + receiver.abort_notified = True + + receiver.arm_abort_intent = arm_abort_intent return SimpleNamespace( - req=SimpleNamespace(bootstrap_room=room), + req=SimpleNamespace( + bootstrap_room=room, + rid=f"request-{room}", + return_logprob=False, + ), kv_receiver=receiver, metadata_buffer_index=idx, ) @@ -159,10 +179,7 @@ def test_holds_until_drained_then_releases(self): mgr = _make_manager() room, idx = 200, 7 q = _make_queue() - dreq = _make_decode_req(room, idx, mgr, n_prefill_ranks=2) - # In production the room is armed in abort_request when the ABORT is - # sent, before the scheduler defers here. - mgr.register_deferred_abort_room(room) + dreq = _make_decode_req(room, idx, mgr, ranks=(4, 9)) q._defer_release(dreq) with patch.object(decode_mod, "release_kv_cache") as rel: @@ -172,13 +189,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, GEN_A, 4) 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, GEN_A, 9) q.resolve_deferred_releases() rel.assert_called_once_with(dreq.req, q.tree_cache, is_insert=False) @@ -186,35 +203,36 @@ def test_holds_until_drained_then_releases(self): self.assertEqual(q._deferred_releases, []) self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [idx]) self.assertEqual(q.metadata_buffers.bootstrap_room[idx], 0) - self.assertNotIn(room, mgr._deferred_abort_ack_tracker) + self.assertNotIn((room, GEN_A), mgr._deferred_abort_ack_tracker) self.assertIsNone(dreq.kv_receiver) - def test_releases_on_timeout_without_ack(self): + def test_timeout_only_diagnoses_retries_and_retains_quarantine(self): mgr = _make_manager() room, idx = 300, 3 q = _make_queue(timeout=30.0) - dreq = _make_decode_req(room, idx, mgr, n_prefill_ranks=1) - # Force an already-expired deadline (no ack will ever arrive). - q._deferred_releases.append((dreq, float("-inf"), idx, 1)) + dreq = _make_decode_req(room, idx, mgr) + q._defer_release(dreq) + q._deferred_releases[0].next_diagnostic_at = float("-inf") with patch.object(decode_mod, "release_kv_cache") as rel: q.resolve_deferred_releases() - rel.assert_called_once_with(dreq.req, q.tree_cache, is_insert=False) + rel.assert_not_called() - self.assertEqual(q._deferred_releases, []) - self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [idx]) - self.assertIsNone(dreq.kv_receiver) + self.assertEqual(len(q._deferred_releases), 1) + self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, []) + self.assertIs(dreq.kv_receiver, q._deferred_releases[0].kv_receiver) + dreq.kv_receiver.retry_abort_notification.assert_called_once_with() + self.assertEqual(q._deferred_releases[0].diagnostic_count, 1) - def test_failed_release_is_isolated_and_not_retried(self): - # A raising _do_release must drop the entry (no double-free on retry) and - # not brick resolve for the remaining entries or subsequent calls. + def test_failed_release_is_fail_stop_while_other_owner_completes(self): mgr = _make_manager() q = _make_queue() - good = _make_decode_req(700, 1, mgr) - bad = _make_decode_req(701, 2, mgr) - # Both already past deadline -> both selected for release. - q._deferred_releases.append((bad, float("-inf"), 2, 1)) - q._deferred_releases.append((good, float("-inf"), 1, 1)) + bad = _make_decode_req(701, 2, mgr, generation=GEN_A) + good = _make_decode_req(700, 1, mgr, generation=GEN_B) + q._defer_release(bad) + q._defer_release(good) + mgr.note_abort_ack(701, GEN_A, 0) + mgr.note_abort_ack(700, GEN_B, 0) calls = [] @@ -224,12 +242,97 @@ def fake_release(req, tree_cache, is_insert): raise RuntimeError("boom") with patch.object(decode_mod, "release_kv_cache", side_effect=fake_release): - q.resolve_deferred_releases() # must not raise - # The good one still released despite the bad one throwing. + q.resolve_deferred_releases() self.assertIn(good.req, calls) - # Nothing left held, and a second call is a clean no-op (no retry). - self.assertEqual(q._deferred_releases, []) + self.assertEqual(len(q._deferred_releases), 1) + failed = q._deferred_releases[0] + self.assertIs(failed.decode_req, bad) + self.assertEqual(failed.release_phase, "KV_RELEASE_ENTERED") + self.assertIsNotNone(failed.release_error) + prior_calls = list(calls) q.resolve_deferred_releases() + self.assertEqual(calls, prior_calls) + + def test_lost_return_from_non_idempotent_free_is_never_retried(self): + mgr = _make_manager() + q = _make_queue() + dreq = _make_decode_req(702, 5, mgr) + q._defer_release(dreq) + mgr.note_abort_ack(702, GEN_A, 0) + + real_free = q.req_to_metadata_buffer_idx_allocator.free + + def committed_then_interrupted(idx): + real_free(idx) + raise KeyboardInterrupt("lost return") + + q.req_to_metadata_buffer_idx_allocator.free = MagicMock( + side_effect=committed_then_interrupted + ) + with self.assertRaises(KeyboardInterrupt): + q.resolve_deferred_releases() + self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [5]) + self.assertEqual( + q._deferred_releases[0].release_phase, + "METADATA_INDEX_FREE_ENTERED", + ) + q.resolve_deferred_releases() + q.req_to_metadata_buffer_idx_allocator.free.assert_called_once_with(5) + + def test_explicit_false_environment_cannot_authorize_timeout_release(self): + mgr = _make_manager() + with patch.dict( + os.environ, + {"SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE": "false"}, + ): + q = DecodeTransferQueue( + gloo_group=None, + req_to_metadata_buffer_idx_allocator=_FakeIdxAllocator(), + tp_rank=0, + metadata_buffers=SimpleNamespace(bootstrap_room={}), + scheduler=SimpleNamespace(spec_algorithm=None), + tree_cache=object(), + ) + self.assertTrue(q.enable_deferred_kv_release) + dreq = _make_decode_req(703, 6, mgr) + q._defer_release(dreq) + q._deferred_releases[0].next_diagnostic_at = float("-inf") + with patch.object(decode_mod, "release_kv_cache") as rel: + q.resolve_deferred_releases() + rel.assert_not_called() + self.assertEqual(len(q._deferred_releases), 1) + + def test_false_compatibility_fields_cannot_bypass_failure_quarantine(self): + mgr = _make_manager() + mgr.enable_deferred_decode_kv_release = False + q = _make_queue() + q.enable_deferred_kv_release = False + q.enable_staging = False + q.gloo_group = None + q.tp_rank = 0 + q.queue = [_make_decode_req(704, 8, mgr)] + q.queue[0].hicache_restore_status = None + # This is specifically a decode-initiated failure with a possible + # remote writer. A propagated prefill failure has already stopped its + # writer and intentionally uses the immediate-safe release path. + q.queue[0].kv_receiver.arm_abort_intent() + q._poll_with_metadata_gate = MagicMock(return_value=[decode_mod.KVPoll.Failed]) + q._clean_hicache_prefetch_resources = MagicMock() + q.scheduler = SimpleNamespace( + enable_decode_hicache=False, + enable_hisparse=False, + output_streamer=SimpleNamespace(stream_output=MagicMock()), + metrics_reporter=SimpleNamespace(enable_metrics=False), + ) + + with patch.object(decode_mod, "release_kv_cache") as release: + q.pop_transferred() + + release.assert_not_called() + self.assertEqual(q.queue, []) + self.assertEqual(len(q._deferred_releases), 1) + self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, []) + self.assertIsNotNone(q._deferred_releases[0].decode_req.kv_receiver) def test_defer_release_records_deadline_and_idx(self): mgr = _make_manager() @@ -237,10 +340,73 @@ def test_defer_release_records_deadline_and_idx(self): dreq = _make_decode_req(room=400, idx=9, mgr=mgr) q._defer_release(dreq) self.assertEqual(len(q._deferred_releases), 1) - held_req, deadline, held_idx, required = q._deferred_releases[0] - self.assertIs(held_req, dreq) - self.assertEqual(held_idx, 9) - self.assertIsInstance(deadline, float) + held = q._deferred_releases[0] + self.assertIs(held.decode_req, dreq) + self.assertEqual(held.metadata_buffer_index, 9) + self.assertEqual(held.transfer_generation, GEN_A) + self.assertIsInstance(held.next_diagnostic_at, float) + + +class TestTransferGenerationPlumbing(CustomTestCase): + @staticmethod + def _make_prealloc_queue(): + queue = DecodePreallocQueue.__new__(DecodePreallocQueue) + queue.transfer_backend = TransferBackend.NIXL + queue.scheduler = SimpleNamespace( + output_streamer=SimpleNamespace(stream_output=MagicMock()) + ) + queue.pending_reqs = [] + queue.retracted_queue = [] + queue._check_if_req_exceed_kv_capacity = MagicMock(return_value=False) + queue._resolve_prefill_dp_rank = MagicMock(return_value=0) + receiver = SimpleNamespace(init=MagicMock()) + queue._create_receiver_and_enqueue = MagicMock( + return_value=SimpleNamespace(kv_receiver=receiver) + ) + return queue, receiver + + def test_rebootstrap_advances_identically_on_each_decode_rank(self): + advanced = [] + for _ in range(2): + queue, receiver = self._make_prealloc_queue() + req = SimpleNamespace( + bootstrap_host="prefill.internal", + transfer_generation=GEN_A, + return_logprob=False, + ) + + queue.add(req, is_rebootstrap=True) + + advanced.append(req.transfer_generation) + queue._create_receiver_and_enqueue.assert_called_once_with( + req, is_rebootstrap=True + ) + receiver.init.assert_called_once_with(0) + + self.assertEqual(advanced, [next_transfer_generation(GEN_A)] * 2) + self.assertNotIn(advanced[0], (GEN_A, GEN_B)) + self.assertEqual(len(advanced[0]), 32) + + def test_missing_generation_fails_request_before_receiver_construction(self): + queue, _ = self._make_prealloc_queue() + req = SimpleNamespace( + bootstrap_host="prefill.internal", + transfer_generation=None, + return_logprob=False, + ) + + with patch.object(decode_mod, "prepare_abort") as prepare_abort: + queue.add(req) + + queue._create_receiver_and_enqueue.assert_not_called() + prepare_abort.assert_called_once() + self.assertEqual( + prepare_abort.call_args.kwargs["status_code"], + decode_mod.HTTPStatus.INTERNAL_SERVER_ERROR, + ) + queue.scheduler.output_streamer.stream_output.assert_called_once_with( + [req], False + ) if __name__ == "__main__": diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index a2bdf5aabdb8..39d5d17373d9 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -14,6 +14,8 @@ ) from sglang.srt.disaggregation.common.staging_handler import ( DecodeStagingHandler, + STAGING_REQ_WIRE_TAG, + STAGING_RSP_WIRE_TAG, handle_staging_req, ) from sglang.srt.disaggregation.common.utils import ( @@ -258,9 +260,11 @@ def test_new_watermark_subscriber_receives_current_allocator_state(self): class TestMooncakePPStaging(unittest.TestCase): def test_staging_response_targets_requesting_pp_rank(self): + generation = "0123456789abcdef0123456789abcdef" sock = Mock() receiver = SimpleNamespace( chunk_staging_infos=[], + transfer_generation=generation, _connect_to_bootstrap_server=Mock(return_value=(sock, threading.Lock())), ) allocator = SimpleNamespace( @@ -275,7 +279,15 @@ def test_staging_response_targets_requesting_pp_rank(self): target = {"pp_rank": 3} handle_staging_req( - [b"STAGING_REQ", b"7", b"0", b"1", b"peer", b"3"], + [ + STAGING_REQ_WIRE_TAG, + generation.encode("ascii"), + b"7", + b"0", + b"1", + b"peer", + b"3", + ], allocator, kv_args, attn_tp_size=16, @@ -287,6 +299,9 @@ def test_staging_response_targets_requesting_pp_rank(self): receiver._connect_to_bootstrap_server.assert_called_once_with(target) sock.send_multipart.assert_called_once() + response = sock.send_multipart.call_args.args[0] + self.assertEqual(response[0], STAGING_RSP_WIRE_TAG) + self.assertEqual(response[1], generation.encode("ascii")) @patch( "sglang.srt.disaggregation.common.staging_buffer.gather_all_layers_to_staging" diff --git a/test/registered/unit/disaggregation/test_mooncake_generation_control.py b/test/registered/unit/disaggregation/test_mooncake_generation_control.py new file mode 100644 index 000000000000..4180863552e2 --- /dev/null +++ b/test/registered/unit/disaggregation/test_mooncake_generation_control.py @@ -0,0 +1,192 @@ +"""Adversarial Mooncake control-plane lease identity tests.""" + +import struct +import threading +from collections import defaultdict +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from sglang.srt.disaggregation.base.conn import KVPoll +from sglang.srt.disaggregation.common.conn import CommonKVReceiver +from sglang.srt.disaggregation.common.staging_handler import DecodeStagingHandler +from sglang.srt.disaggregation.mooncake.conn import ( + MOONCAKE_AUX_DATA_WIRE_TAG, + MOONCAKE_CHUNK_READY_WIRE_TAG, + MOONCAKE_TRANSFER_STATUS_WIRE_TAG, + MooncakeKVManager, + MooncakeKVReceiver, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + +GEN_A = "a1" * 16 +GEN_B = "b2" * 16 +ROOM = 7 +SESSION = "decode-session" + + +def _manager(generation=GEN_B, expected=frozenset({3, 5})): + manager = object.__new__(MooncakeKVManager) + manager._abort_state_lock = threading.Lock() + manager._receive_transfer_leases = {} + manager.prefill_response_tracker = defaultdict(set) + manager.request_status = {ROOM: KVPoll.WaitingForInput} + manager.required_prefill_response_num_table = {ROOM: len(expected)} + manager.failure_records = {} + manager.failure_lock = threading.Lock() + manager.enable_staging = True + manager._staging_handler = MagicMock() + manager.kv_args = SimpleNamespace() + manager._register_receive_transfer_lease(ROOM, generation, SESSION, expected) + return manager + + +def _chunk(generation, rank=3, session=SESSION): + return [ + MOONCAKE_CHUNK_READY_WIRE_TAG, + generation.encode("ascii"), + str(ROOM).encode("ascii"), + b"0", + b"4", + b"2", + session.encode("ascii"), + str(rank).encode("ascii"), + ] + + +def _aux(generation, rank=3, payload=b"data"): + return [ + MOONCAKE_AUX_DATA_WIRE_TAG, + generation.encode("ascii"), + str(ROOM).encode("ascii"), + str(rank).encode("ascii"), + b"0", + b"1", + struct.pack(">I", len(payload)), + payload, + ] + + +def _status(generation, rank, status=KVPoll.Success): + return [ + MOONCAKE_TRANSFER_STATUS_WIRE_TAG, + generation.encode("ascii"), + str(ROOM).encode("ascii"), + str(status).encode("ascii"), + str(rank).encode("ascii"), + ] + + +class TestMooncakeGenerationControl(CustomTestCase): + def test_receiver_init_binds_immutable_exact_rank_set(self): + manager = _manager() + manager._clear_receive_transfer_lease(ROOM, GEN_B) + manager.addr_to_rooms_tracker = defaultdict(set) + manager.get_session_id = MagicMock(return_value=SESSION) + + receiver = MooncakeKVReceiver( + manager, + "prefill:8998", + ROOM, + transfer_generation=GEN_B, + ) + receiver.bootstrap_infos = [ + {"prefill_unique_rank": 3, "is_dummy": False}, + {"prefill_unique_rank": 4, "is_dummy": True}, + {"prefill_unique_rank": 5, "is_dummy": False}, + ] + with patch.object(CommonKVReceiver, "init", return_value=None): + receiver.init(prefill_dp_rank=0) + + lease = manager._receive_transfer_leases[ROOM] + self.assertEqual(lease.generation, GEN_B) + self.assertEqual(lease.session_id, SESSION) + self.assertEqual(lease.expected_prefill_ranks, frozenset({3, 5})) + with self.assertRaisesRegex(RuntimeError, "immutable"): + manager._register_receive_transfer_lease( + ROOM, GEN_B, SESSION, frozenset({3, 6}) + ) + + def test_old_chunk_cannot_scatter_reused_room(self): + manager = _manager() + + self.assertFalse(manager._handle_chunk_ready(_chunk(GEN_A))) + self.assertFalse(manager._handle_chunk_ready(_chunk(GEN_B, rank=99))) + self.assertFalse(manager._handle_chunk_ready(_chunk(GEN_B, session="old"))) + self.assertFalse( + manager._handle_chunk_ready( + [b"CHUNK_READY", b"7", b"0", b"4", b"2", SESSION.encode()] + ) + ) + manager._staging_handler.handle_chunk_arrived.assert_not_called() + + self.assertTrue(manager._handle_chunk_ready(_chunk(GEN_B, rank=3))) + manager._staging_handler.handle_chunk_arrived.assert_called_once_with( + ROOM, 0, 4, 2, 3 + ) + + def test_old_aux_cannot_write_reused_room(self): + manager = _manager() + target = "sglang.srt.disaggregation.mooncake.conn.AuxDataCodec.deserialize_data_to_buffer" + with patch(target) as deserialize: + self.assertFalse(manager._handle_aux_data(_aux(GEN_A))) + self.assertFalse(manager._handle_aux_data(_aux(GEN_B, rank=99))) + self.assertFalse( + manager._handle_aux_data( + [b"AUX_DATA", b"7", b"0", b"1", struct.pack(">I", 4), b"data"] + ) + ) + deserialize.assert_not_called() + + self.assertTrue(manager._handle_aux_data(_aux(GEN_B, rank=3))) + deserialize.assert_called_once_with(manager.kv_args, 0, 1, b"data") + + def test_status_requires_exact_generation_and_all_distinct_expected_ranks(self): + manager = _manager() + + self.assertFalse(manager._handle_transfer_status(_status(GEN_A, 3))) + self.assertFalse(manager._handle_transfer_status(_status(GEN_B, 99))) + self.assertFalse(manager._handle_transfer_status([b"7", b"4", b"3"])) + self.assertEqual(manager.request_status[ROOM], KVPoll.WaitingForInput) + + self.assertTrue(manager._handle_transfer_status(_status(GEN_B, 3))) + self.assertTrue(manager._handle_transfer_status(_status(GEN_B, 3))) + self.assertEqual(manager.request_status[ROOM], KVPoll.WaitingForInput) + self.assertEqual(manager.prefill_response_tracker[ROOM], {3}) + + self.assertTrue(manager._handle_transfer_status(_status(GEN_B, 5))) + self.assertEqual(manager.prefill_response_tracker[ROOM], {3, 5}) + self.assertEqual(manager.request_status[ROOM], KVPoll.Success) + + def test_delayed_old_receiver_clear_cannot_erase_new_lease(self): + manager = _manager(generation=GEN_A) + self.assertTrue(manager._clear_receive_transfer_lease(ROOM, GEN_A)) + manager.request_status[ROOM] = KVPoll.WaitingForInput + manager.required_prefill_response_num_table[ROOM] = 2 + # The old receiver has already cleared its identity. A duplicate clear + # must not treat that absence as authority over state B is publishing. + self.assertFalse(manager._clear_receive_transfer_lease(ROOM, GEN_A)) + self.assertEqual(manager.request_status[ROOM], KVPoll.WaitingForInput) + manager._register_receive_transfer_lease( + ROOM, GEN_B, SESSION, frozenset({3, 5}) + ) + + self.assertFalse(manager._clear_receive_transfer_lease(ROOM, GEN_A)) + self.assertEqual(manager._receive_transfer_leases[ROOM].generation, GEN_B) + self.assertEqual(manager.request_status[ROOM], KVPoll.WaitingForInput) + + def test_staging_fan_in_deduplicates_exact_writer_rank(self): + handler = object.__new__(DecodeStagingHandler) + handler._room_to_receiver = {ROOM: object()} + handler._writer_counts = {} + handler.num_writers_for = MagicMock(return_value=2) + handler.submit_chunk_scatter = MagicMock(return_value=True) + + self.assertFalse(handler.handle_chunk_arrived(ROOM, 0, 4, 2, 3)) + self.assertFalse(handler.handle_chunk_arrived(ROOM, 0, 4, 2, 3)) + handler.submit_chunk_scatter.assert_not_called() + + self.assertTrue(handler.handle_chunk_arrived(ROOM, 0, 4, 2, 5)) + handler.submit_chunk_scatter.assert_called_once_with(ROOM, 0, 4, 2) diff --git a/test/registered/unit/disaggregation/test_nixl_backend_basic.py b/test/registered/unit/disaggregation/test_nixl_backend_basic.py index 2bc8a210738a..e6ab6c0c3c11 100644 --- a/test/registered/unit/disaggregation/test_nixl_backend_basic.py +++ b/test/registered/unit/disaggregation/test_nixl_backend_basic.py @@ -6,6 +6,7 @@ import types import unittest from collections import defaultdict +from dataclasses import replace from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -13,7 +14,13 @@ from sglang.srt.disaggregation.base.conn import KVPoll from sglang.srt.disaggregation.common.conn import CommonKVManager -from sglang.srt.disaggregation.common.staging_handler import PrefillStagingContext +from sglang.srt.disaggregation.common.staging_handler import ( + STAGING_REQ_WIRE_TAG, + STAGING_RSP_WIRE_TAG, + PrefillStagingContext, + handle_staging_req, + handle_staging_rsp, +) from sglang.srt.disaggregation.common.utils import pack_int_lists from sglang.srt.disaggregation.nixl.conn import ( KVArgsRegisterInfo, @@ -23,12 +30,21 @@ TransferInfo, TransferKVChunk, TransferStatus, + repeat_indices_over_layers, ) +from sglang.srt.disaggregation.utils import DisaggregationMode 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") +TRANSFER_GENERATION = "11" * 16 +OLD_TRANSFER_GENERATION = "22" * 16 + + +def completion_notification(room, suffix, generation=TRANSFER_GENERATION): + return f"nixlv2_{generation}_{room}_{suffix}" + class NotificationFakeAgent: def __init__(self, messages): @@ -120,6 +136,8 @@ def test_from_zmq_parses_required_fields(self): b"2", pack_int_lists(state_indices, "i"), b"11", + b"0", + TRANSFER_GENERATION.encode("ascii"), ] info = TransferInfo.from_zmq(msg) @@ -133,6 +151,7 @@ def test_from_zmq_parses_required_fields(self): self.assertEqual(info.required_dst_info_num, 2) self.assertEqual(info.dst_state_indices, state_indices) self.assertEqual(info.decode_prefix_len, 11) + self.assertEqual(info.transfer_generation, TRANSFER_GENERATION) def test_from_zmq_defaults_optional_fields(self): info = TransferInfo.from_zmq( @@ -144,6 +163,10 @@ def test_from_zmq_defaults_optional_fields(self): np.array([1], dtype=np.int32).tobytes(), b"0", b"1", + b"", + b"", + b"", + TRANSFER_GENERATION.encode("ascii"), ] ) @@ -162,6 +185,8 @@ def test_decode_radix_full_hit_is_not_dummy(self): b"1", b"", b"128", + b"", + TRANSFER_GENERATION.encode("ascii"), ] ) @@ -179,6 +204,8 @@ def test_empty_indices_without_decode_prefix_is_dummy(self): b"1", b"", b"0", + b"", + TRANSFER_GENERATION.encode("ascii"), ] ) @@ -276,8 +303,20 @@ def test_from_zmq_allows_missing_state_and_staging_fields(self): class TestNixlTransferStatus(CustomTestCase): + def test_completion_source_rank_set_is_exact_and_immutable(self): + status = TransferStatus(transfer_generation=TRANSFER_GENERATION) + status.bind_expected_source_ranks(frozenset({4, 9})) + + self.assertTrue(status.accepts_source_rank(4)) + for unexpected in (True, -1, 0, 10): + self.assertFalse(status.accepts_source_rank(unexpected)) + self.assertEqual(status.num_pp_ranks_expected, 2) + status.bind_expected_source_ranks(frozenset({4, 9})) + with self.assertRaises(RuntimeError): + status.bind_expected_source_ranks(frozenset({4})) + def test_not_done_until_aux_and_expected_count_arrive(self): - status = TransferStatus() + status = TransferStatus(transfer_generation=TRANSFER_GENERATION) self.assertFalse(status.is_done()) @@ -294,7 +333,7 @@ def test_not_done_until_aux_and_expected_count_arrive(self): self.assertTrue(status.is_done()) def test_zero_kv_aux_only_completion(self): - status = TransferStatus() + status = TransferStatus(transfer_generation=TRANSFER_GENERATION) status.received_aux = True status.num_pp_ranks_expected = 1 status.expected_kvs_per_pp[0] = 0 @@ -302,7 +341,7 @@ def test_zero_kv_aux_only_completion(self): self.assertTrue(status.is_done()) def test_multi_pp_requires_each_rank_expected_chunks(self): - status = TransferStatus() + status = TransferStatus(transfer_generation=TRANSFER_GENERATION) status.received_aux = True status.num_pp_ranks_expected = 2 status.expected_kvs_per_pp[0] = 1 @@ -315,7 +354,7 @@ def test_multi_pp_requires_each_rank_expected_chunks(self): self.assertTrue(status.is_done()) def test_state_required_completion_waits_for_all_pp_ranks(self): - status = TransferStatus() + status = TransferStatus(transfer_generation=TRANSFER_GENERATION) status.received_aux = True status.num_pp_ranks_expected = 2 status.expected_kvs_per_pp[0] = 0 @@ -339,6 +378,148 @@ def test_last_zero_page_chunk_is_sent_for_aux_only_completion(self): self.assertFalse(sender.should_send_kv_chunk(0, last_chunk=False)) self.assertTrue(sender.should_send_kv_chunk(3, last_chunk=False)) + def test_both_data_planes_support_overlap_early_send(self): + sender = object.__new__(NixlKVSender) + + for use_torch_transfer in (False, True): + sender.kv_mgr = SimpleNamespace(_use_torch_transfer=use_torch_transfer) + self.assertTrue(sender.supports_overlap_early_send) + + def test_send_consumes_early_send_event_into_transfer_chunk(self): + sender = object.__new__(NixlKVSender) + wait_event = object() + sender._early_send_wait_event = wait_event + sender._send_failed = False + sender._transfer_start_time = 1.0 + sender.bootstrap_room = 7 + sender.chunk_id = 2 + sender.aux_index = 3 + sender.has_sent = False + sender.kv_mgr = SimpleNamespace(add_transfer_request=MagicMock()) + prepared_indices = np.array([5], dtype=np.int32) + sender._prepare_send_indices = MagicMock( + return_value=( + prepared_indices, + slice(0, 1), + False, + False, + ) + ) + sender._record_transfer_indices = MagicMock() + + sender.send(np.array([5], dtype=np.int32), num_kv_tokens=4) + + self.assertIsNone(sender._early_send_wait_event) + sender.kv_mgr.add_transfer_request.assert_called_once_with( + 7, + prepared_indices, + slice(0, 1), + False, + 2, + 3, + None, + 4, + wait_event=wait_event, + ) + + +class TestNixlBootstrapErrors(CustomTestCase): + def _make_manager(self): + mgr = object.__new__(NixlKVManager) + mgr.enable_staging = False + mgr._bootstrap_errors = {} + mgr.exceptions = {} + mgr.transfer_infos = {} + mgr.req_to_decode_prefix_len = {} + mgr.prep_handles = {} + mgr._handle_abort_notification = MagicMock(return_value=False) + mgr.record_failure = MagicMock() + mgr.update_status = MagicMock() + return mgr + + def test_registration_error_is_propagated_to_later_room(self): + mgr = self._make_manager() + unsupported = NotImplementedError("unsupported PyTorch transfer topology") + mgr._add_remote_peer = MagicMock(side_effect=unsupported) + registration_message = [ + b"NixlMsgGuard", + b"None", + b"endpoint", + b"1234", + b"decode-agent", + ] + room_message = [ + b"NixlMsgGuard", + b"41", + b"endpoint", + b"1234", + b"decode-agent", + b"", + b"0", + b"1", + b"", + b"", + b"0", + TRANSFER_GENERATION.encode("ascii"), + ] + + with patch.object(KVArgsRegisterInfo, "from_zmq", return_value=object()): + with self.assertRaises(NotImplementedError): + mgr._handle_bootstrap_message(registration_message) + self.assertIs(mgr._bootstrap_errors["decode-agent"], unsupported) + + with self.assertRaises(NotImplementedError): + mgr._handle_bootstrap_message(room_message) + + self.assertIs(mgr.exceptions[41], unsupported) + mgr.record_failure.assert_called_once_with(41, str(unsupported)) + mgr.update_status.assert_called_once_with(41, KVPoll.Failed) + + def test_listener_continues_after_message_error(self): + mgr = self._make_manager() + mgr.server_socket = SimpleNamespace( + recv_multipart=MagicMock(side_effect=[[b"bad"], [b"good"], SystemExit()]) + ) + mgr._handle_bootstrap_message = MagicMock( + side_effect=[ValueError("bad bootstrap"), None] + ) + + with patch( + "sglang.srt.disaggregation.nixl.conn.threading.Thread" + ) as thread_cls: + mgr._start_bootstrap_thread() + target = thread_cls.call_args.kwargs["target"] + + with self.assertRaises(SystemExit): + target() + + self.assertEqual(mgr._handle_bootstrap_message.call_count, 2) + self.assertTrue(thread_cls.call_args.kwargs["daemon"]) + + def test_peer_import_failure_does_not_poison_registration_table(self): + mgr = self._make_manager() + mgr._use_torch_transfer = True + mgr.attn_tp_size = 1 + mgr.disaggregation_mode = object() + mgr.decode_kv_args_table = {} + mgr.requires_dcp_relayout = MagicMock(return_value=False) + mgr.agent = SimpleNamespace( + add_remote_agent=MagicMock(side_effect=RuntimeError("metadata rejected")) + ) + peer = SimpleNamespace( + agent_name="decode-agent", + agent_metadata=b"metadata", + decode_tp_size=1, + dst_dcp_size=1, + dst_dcp_rank=0, + dst_kv_mem_kinds=["VRAM"], + ) + + with self.assertRaisesRegex(RuntimeError, "metadata rejected"): + mgr._add_remote_peer(peer) + + self.assertNotIn("decode-agent", mgr.decode_kv_args_table) + class TestNixlAbortHandling(CustomTestCase): def _make_manager(self, request_status=None): @@ -347,9 +528,15 @@ def _make_manager(self, request_status=None): mgr._connect = MagicMock() mgr.failure_lock = threading.Lock() mgr.failure_records = {} - # These cases cover the legacy no-ack behavior; the deferred-release ack - # path is exercised in test_nixl_deferred_kv_release.py. + mgr._abort_state_lock = threading.Lock() + mgr._deferred_ack_targets = {} + mgr._active_transfer_generations = {} + mgr._abort_ack_receipts = set() + mgr._staging_outstanding = {room: 1 for room in mgr.request_status} + mgr._send_abort_ack = MagicMock(return_value=True) mgr.enable_deferred_decode_kv_release = False + for room in mgr.request_status: + mgr.register_transfer_generation(room, TRANSFER_GENERATION) return mgr def test_given_known_incomplete_room_when_abort_arrives_then_room_fails_without_ack( @@ -358,7 +545,13 @@ def test_given_known_incomplete_room_when_abort_arrives_then_room_fails_without_ mgr = self._make_manager({11: KVPoll.WaitingForInput}) handled = mgr._handle_abort_notification( - [b"ABORT", b"11", b"127.0.0.1", b"5555"] + [ + b"ABORT", + b"11", + b"127.0.0.1", + b"5555", + TRANSFER_GENERATION.encode("ascii"), + ] ) self.assertTrue(handled) @@ -373,7 +566,13 @@ def test_given_successful_room_when_abort_arrives_then_status_is_preserved(self) mgr = self._make_manager({12: KVPoll.Success}) handled = mgr._handle_abort_notification( - [b"ABORT", b"12", b"127.0.0.1", b"5556"] + [ + b"ABORT", + b"12", + b"127.0.0.1", + b"5556", + TRANSFER_GENERATION.encode("ascii"), + ] ) self.assertTrue(handled) @@ -385,7 +584,13 @@ def test_given_unknown_room_when_abort_arrives_then_status_remains_absent(self): mgr = self._make_manager() handled = mgr._handle_abort_notification( - [b"ABORT", b"14", b"127.0.0.1", b"5557"] + [ + b"ABORT", + b"14", + b"127.0.0.1", + b"5557", + TRANSFER_GENERATION.encode("ascii"), + ] ) self.assertTrue(handled) @@ -397,7 +602,13 @@ def test_given_malformed_abort_when_handled_then_no_exception_or_ack(self): mgr = self._make_manager({13: KVPoll.WaitingForInput}) handled = mgr._handle_abort_notification( - [b"ABORT", b"invalid-room", b"127.0.0.1", b"5558"] + [ + b"ABORT", + b"invalid-room", + b"127.0.0.1", + b"5558", + TRANSFER_GENERATION.encode("ascii"), + ] ) self.assertTrue(handled) @@ -446,6 +657,7 @@ def _make_manager(self, room): dst_aux_index=0, required_dst_info_num=1, dst_state_indices=[], + transfer_generation=TRANSFER_GENERATION, ) } } @@ -472,11 +684,16 @@ def _make_manager(self, room): mgr.enable_deferred_decode_kv_release = False mgr._staging_ctx = None mgr._staging_outstanding = defaultdict(int) + mgr._abort_state_lock = threading.Lock() + mgr._deferred_ack_targets = {} + mgr._active_transfer_generations = {room: TRANSFER_GENERATION} + mgr._abort_ack_receipts = set() mgr.is_mla_backend = False mgr.is_hybrid_mla_backend = False mgr.attn_tp_size = 1 mgr.transfer_source_rank = 0 mgr.kv_args = SimpleNamespace(engine_rank=0, kv_data_ptrs=[0]) + mgr._use_torch_transfer = False mgr.exceptions = {} mgr.failure_lock = threading.Lock() mgr.failure_records = {} @@ -518,9 +735,12 @@ def test_given_last_chunk_aborts_mid_transfer_when_worker_finishes_then_failed_s self.assertNotIn(room, mgr.transfer_infos) self.assertNotIn(room, mgr.req_to_decode_prefix_len) mgr.send_aux.assert_called_once() - self.assertEqual(mgr.send_aux.call_args.args[-1], "21_aux_nokv_0_0") + self.assertEqual( + mgr.send_aux.call_args.args[-1], + completion_notification(21, "aux_nokv_0_0"), + ) - def test_given_non_last_chunk_aborts_mid_transfer_when_worker_finishes_then_failed_status_is_preserved( + def test_given_non_last_chunk_aborts_mid_transfer_then_failed_room_is_retired( self, ): room = 22 @@ -531,12 +751,247 @@ def test_given_non_last_chunk_aborts_mid_transfer_when_worker_finishes_then_fail self._run_worker_once(mgr, chunk) self.assertEqual(mgr.request_status[room], KVPoll.Failed) - self.assertIn(room, mgr.transfer_infos) - self.assertIn(room, mgr.req_to_decode_prefix_len) + self.assertNotIn(room, mgr.transfer_infos) + self.assertNotIn(room, mgr.req_to_decode_prefix_len) + self.assertIn((room, TRANSFER_GENERATION), mgr._abort_ack_receipts) mgr.send_kvcache.assert_called_once() + def test_waits_for_early_send_event_once_before_source_access(self): + for room, use_torch_transfer in ((23, False), (24, True)): + with self.subTest(use_torch_transfer=use_torch_transfer): + mgr = self._make_manager(room) + mgr._use_torch_transfer = use_torch_transfer + mgr.agent.progress = MagicMock() + actions = [] + captured_handles = [] + mgr.agent.begin_handle_batch = lambda handles: captured_handles.append( + handles + ) + mgr.agent.end_handle_batch = lambda handles: self.assertIs( + captured_handles[-1], handles + ) + wait_event = MagicMock() + wait_event.synchronize.side_effect = lambda: actions.append("wait") + + def send_kvcache(*args, **kwargs): + actions.append("send") + if use_torch_transfer: + captured_handles[-1].append("handle") + return "handle" + + mgr.send_kvcache = MagicMock(side_effect=send_kvcache) + chunk = self._make_chunk(room, [1], is_last_chunk=False) + chunk.wait_event = wait_event + + self._run_worker_once(mgr, chunk) + + self.assertEqual(actions, ["wait", "send"]) + wait_event.synchronize.assert_called_once_with() + self.assertIsNone(chunk.wait_event) + if use_torch_transfer: + mgr.agent.progress.assert_called_once_with() + else: + mgr.agent.progress.assert_not_called() + + def test_torch_failure_is_published_only_after_every_handle_is_drained(self): + room = 26 + mgr = self._make_manager(room) + mgr._use_torch_transfer = True + second_req = replace(mgr.transfer_infos[room]["agent"], agent_name="agent-2") + mgr.transfer_infos[room]["agent-2"] = second_req + mgr.decode_kv_args_table["agent-2"] = mgr.decode_kv_args_table["agent"] + release_kv_cache = MagicMock() + events = [] + batch = [] + original_update_status = mgr.update_status + + def send_kvcache(*args, **kwargs): + handle = "failed" if not batch[-1] else "active" + batch[-1].append(handle) + return handle + + def update_status(candidate_room, status): + events.append(("status", status)) + original_update_status(candidate_room, status) + if status == KVPoll.Failed: + release_kv_cache(candidate_room) + + def cancel_handles_and_wait(handles): + events.append(("drain", tuple(handles))) + self.assertEqual(mgr.request_status[room], KVPoll.Transferring) + release_kv_cache.assert_not_called() + + mgr.update_status = update_status + mgr.send_kvcache = MagicMock(side_effect=send_kvcache) + mgr.agent = SimpleNamespace( + begin_handle_batch=lambda handles: batch.append(handles), + end_handle_batch=lambda handles: self.assertIs(batch[-1], handles), + progress=MagicMock(), + check_xfer_state=MagicMock(return_value="ERR"), + pop_xfer_error=MagicMock(return_value=None), + cancel_handles_and_wait=cancel_handles_and_wait, + ) + chunk = self._make_chunk(room, [1], is_last_chunk=False) + + self._run_worker_once(mgr, chunk) + + self.assertEqual(events[-2], ("drain", ("failed", "active"))) + self.assertEqual(events[-1], ("status", KVPoll.Failed)) + release_kv_cache.assert_called_once_with(room) + self.assertEqual(mgr._staging_outstanding[room], 0) + self.assertFalse(chunk.staging_counted) + + def test_fatal_worker_cut_waits_without_timeout_until_batch_is_quiescent(self): + class FatalWorkerCut(BaseException): + pass + + room = 28 + mgr = self._make_manager(room) + mgr._use_torch_transfer = True + mgr.enable_deferred_decode_kv_release = True + events = [] + batches = [] + cleanup_attempts = 0 + original_update_status = mgr.update_status + + def begin(handles): + batches.append(handles) + events.append("begin") + + def send_then_die(*args, **kwargs): + batches[-1].append("live-handle") + raise FatalWorkerCut("worker cancelled") + + def retain(handles): + self.assertIs(handles, batches[-1]) + self.assertEqual(mgr.request_status[room], KVPoll.Transferring) + events.append("retain") + + def drain(handles): + nonlocal cleanup_attempts + cleanup_attempts += 1 + events.append(f"drain-{cleanup_attempts}") + self.assertIs(handles, batches[-1]) + self.assertEqual(mgr.request_status[room], KVPoll.Transferring) + if cleanup_attempts == 1: + raise KeyboardInterrupt("cleanup interrupted") + + def update_status(candidate_room, status): + original_update_status(candidate_room, status) + if status == KVPoll.Failed: + events.append("failed") + + def ack(candidate_room, generation=None): + self.assertEqual(candidate_room, room) + self.assertEqual(generation, TRANSFER_GENERATION) + self.assertEqual(mgr.request_status[room], KVPoll.Failed) + self.assertEqual(mgr._staging_outstanding[room], 0) + self.assertFalse(chunk.staging_counted) + events.append("ack") + + def end(handles): + self.assertIs(handles, batches[-1]) + events.append("end") + + mgr.update_status = update_status + mgr.send_kvcache = MagicMock(side_effect=send_then_die) + mgr._maybe_ack_drained_abort = ack + mgr.agent = SimpleNamespace( + begin_handle_batch=begin, + end_handle_batch=end, + retain_handle_batch=retain, + cancel_handles_and_wait=drain, + handle_batch_is_quiescent=lambda handles: False, + ) + chunk = self._make_chunk(room, [1], is_last_chunk=False) + queue = SimpleNamespace(get=MagicMock(return_value=chunk)) + + with self.assertRaisesRegex(FatalWorkerCut, "worker cancelled"): + mgr.transfer_worker(queue) + + self.assertEqual( + events, + ["begin", "retain", "drain-1", "drain-2", "failed", "ack", "end"], + ) + self.assertEqual(mgr._staging_outstanding[room], 0) + self.assertFalse(chunk.staging_counted) + self.assertIsInstance(mgr.exceptions[room], RuntimeError) + + def test_torch_peer_bootstrap_base_exception_rolls_back_exact_state(self): + class BootstrapCut(BaseException): + pass + + class StoreThenInterrupt(dict): + def __setitem__(self, key, value): + super().__setitem__(key, value) + raise BootstrapCut("manager commit cut") + + room = 29 + mgr = self._make_manager(room) + mgr._use_torch_transfer = True + mgr.attn_tp_size = 1 + mgr.disaggregation_mode = DisaggregationMode.PREFILL + mgr.requires_dcp_relayout = MagicMock(return_value=False) + mgr.prep_handles = {"": "existing-local"} + mgr.decode_kv_args_table = StoreThenInterrupt() + released = [] + rolled_back = [] + mgr.agent = SimpleNamespace( + add_remote_agent=MagicMock(return_value="agent"), + release_dlist_handle=lambda handle: released.append(handle), + rollback_remote_agent=lambda name, metadata: rolled_back.append( + (name, metadata) + ), + ) + mgr._prepare_payload_xfer = lambda peer: mgr.prep_handles.__setitem__( + "agent", "new-remote" + ) + peer = SimpleNamespace( + agent_name="agent", + agent_metadata=b"peer-metadata", + decode_tp_size=1, + dst_dcp_size=1, + dst_dcp_rank=0, + dst_kv_mem_kinds=["VRAM"], + ) + + with self.assertRaisesRegex(BootstrapCut, "manager commit cut"): + mgr._add_remote_peer(peer) + + self.assertNotIn("agent", mgr.decode_kv_args_table) + self.assertEqual(mgr.prep_handles, {"": "existing-local"}) + self.assertEqual(released, ["new-remote"]) + self.assertEqual(rolled_back, [("agent", b"peer-metadata")]) + + def test_partial_multi_submit_uses_worker_owned_handle_sink(self): + room = 27 + mgr = self._make_manager(room) + mgr._use_torch_transfer = True + mgr.decode_kv_args_table["agent"].kv_xfer_segments = [object(), object()] + drained = [] + + def submit_then_fail(_peer, _src, _dst, _notif, handles): + handles.append("first-active-part") + raise RuntimeError("second part failed before publication") + + mgr.send_kvcache_mixed = MagicMock(side_effect=submit_then_fail) + mgr.agent = SimpleNamespace( + begin_handle_batch=lambda handles: None, + end_handle_batch=lambda handles: None, + progress=MagicMock(), + check_xfer_state=MagicMock(), + pop_xfer_error=MagicMock(return_value=None), + cancel_handles_and_wait=lambda handles: drained.append(tuple(handles)), + ) + chunk = self._make_chunk(room, [1], is_last_chunk=False) + + self._run_worker_once(mgr, chunk) + + self.assertEqual(drained, [("first-active-part",)]) + self.assertEqual(mgr.request_status[room], KVPoll.Failed) + def test_dcp_destinations_use_disjoint_pack_regions_before_chunk_barrier(self): - room = 23 + room = 25 mgr = self._make_manager(room) agents = ("agent0a", "agent0b", "agent1") dcp_ranks = (0, 0, 1) @@ -550,6 +1005,7 @@ def test_dcp_destinations_use_disjoint_pack_regions_before_chunk_barrier(self): dst_aux_index=0, required_dst_info_num=len(agents), dst_state_indices=[], + transfer_generation=TRANSFER_GENERATION, ) for i, agent in enumerate(agents) } @@ -623,10 +1079,23 @@ def check_xfer_state(_handle): class TestNixlNotifications(CustomTestCase): - def _make_manager(self, messages, required=None): + def _make_manager( + self, + messages, + *, + room=5, + required=None, + generation=TRANSFER_GENERATION, + expected_ranks=frozenset({0}), + ): mgr = object.__new__(NixlKVManager) mgr.agent = NotificationFakeAgent(messages) - mgr.transfer_statuses = defaultdict(TransferStatus) + mgr.transfer_statuses = { + room: TransferStatus( + transfer_generation=generation, + expected_source_ranks=expected_ranks, + ) + } mgr.required_prefill_response_num_table = required or {} mgr.enable_staging = False mgr._staging_handler = None @@ -634,7 +1103,7 @@ def _make_manager(self, messages, required=None): return mgr def test_kv_last_notification_sets_expected_count(self): - mgr = self._make_manager(["5_kv_2_1_0"]) + mgr = self._make_manager([completion_notification(5, "kv_2_1_0")]) mgr.update_transfer_status() @@ -644,19 +1113,26 @@ def test_kv_last_notification_sets_expected_count(self): self.assertEqual(status.num_pp_ranks_expected, 1) def test_staging_notification_preserves_agent_name_with_underscores(self): - mgr = self._make_manager(["5_stg_0_1_0_2_4_8_agent_with_underscores"]) + mgr = self._make_manager( + [completion_notification(5, "stg_0_1_0_2_4_8_agent_with_underscores")] + ) calls = [] mgr._handle_staging_chunk_arrived = lambda *args: calls.append(args) mgr.update_transfer_status() - self.assertEqual(calls, [(5, 2, 4, 8, "agent_with_underscores")]) + self.assertEqual(calls, [(5, 2, 4, 8, 0)]) status = mgr.transfer_statuses[5] self.assertEqual(status.received_kvs_per_pp[0], {0}) self.assertEqual(status.expected_kvs_per_pp[0], 1) def test_aux_nokv_marks_zero_expected_chunks_for_pp_rank(self): - mgr = self._make_manager(["6_aux_nokv_3"], required={6: 4}) + mgr = self._make_manager( + [completion_notification(6, "aux_nokv_3_0")], + room=6, + required={6: 4}, + expected_ranks=frozenset({0, 1, 2, 3}), + ) mgr.update_transfer_status() @@ -666,21 +1142,268 @@ def test_aux_nokv_marks_zero_expected_chunks_for_pp_rank(self): self.assertEqual(status.num_pp_ranks_expected, 4) def test_state_notification_marks_pp_rank(self): - mgr = self._make_manager(["7_state_2"]) + mgr = self._make_manager( + [completion_notification(7, "state_2_4")], + room=7, + expected_ranks=frozenset({2}), + ) mgr.update_transfer_status() self.assertEqual(mgr.transfer_statuses[7].received_state_per_pp, {2}) def test_aux_nokv_allows_full_hit_completion(self): - mgr = self._make_manager(["8_aux_nokv_0"], required={8: 1}) + mgr = self._make_manager( + [completion_notification(8, "aux_nokv_0_0")], + room=8, + required={8: 1}, + ) mgr.update_transfer_status() self.assertTrue(mgr.transfer_statuses[8].is_done()) + def test_mixed_parts_complete_only_after_all_generation_bound_parts(self): + mgr = self._make_manager( + [ + completion_notification(9, "kv_0_1_0_part_0_2"), + completion_notification(9, "kv_0_1_0_part_1_2"), + ], + room=9, + ) + + mgr.update_transfer_status() + + status = mgr.transfer_statuses[9] + self.assertEqual(status.received_kvs_per_pp[0], {0}) + self.assertEqual(status.expected_kvs_per_pp[0], 1) + + def test_arbitrary_and_duplicate_source_ranks_cannot_satisfy_completion(self): + mgr = self._make_manager( + [ + completion_notification(10, "kv_0_1_99"), + completion_notification(10, "aux_99"), + completion_notification(10, "kv_0_1_0"), + completion_notification(10, "kv_0_1_0"), + completion_notification(10, "aux_0"), + ], + room=10, + expected_ranks=frozenset({0, 1}), + ) + + mgr.update_transfer_status() + + status = mgr.transfer_statuses[10] + self.assertEqual(status.received_kvs_per_pp, {0: {0}}) + self.assertEqual(status.expected_kvs_per_pp, {0: 1}) + self.assertTrue(status.received_aux) + self.assertFalse(status.is_done()) + + def test_old_generation_cannot_mutate_any_completion_family_after_reuse(self): + mgr = self._make_manager( + [ + completion_notification(5, "kv_0_1_0", OLD_TRANSFER_GENERATION), + completion_notification( + 5, "kv_0_1_0_part_0_1", OLD_TRANSFER_GENERATION + ), + completion_notification(5, "state_0_0", OLD_TRANSFER_GENERATION), + completion_notification(5, "aux_nokv_0_0", OLD_TRANSFER_GENERATION), + completion_notification( + 5, + "stg_0_1_0_0_0_1_old_agent", + OLD_TRANSFER_GENERATION, + ), + ] + ) + staging_calls = [] + mgr._handle_staging_chunk_arrived = lambda *args: staging_calls.append(args) + + mgr.update_transfer_status() + + status = mgr.transfer_statuses[5] + self.assertEqual(status.received_kvs_per_pp, {}) + self.assertEqual(status.expected_kvs_per_pp, {}) + self.assertEqual(status.received_state_per_pp, set()) + self.assertFalse(status.received_aux) + self.assertIsNone(status.received_kv_parts_per_pp) + self.assertEqual(staging_calls, []) + + def test_legacy_and_malformed_notifications_never_create_or_mutate_status(self): + mgr = self._make_manager( + [ + "5_kv_0_1_0", + f"nixlv1_{TRANSFER_GENERATION}_5_aux", + completion_notification(5, "kv_-1_1_0"), + completion_notification(5, "kv_0_2_0"), + completion_notification(5, "aux_extra"), + completion_notification(999, "aux_0"), + ] + ) + + mgr.update_transfer_status() + + self.assertEqual(set(mgr.transfer_statuses), {5}) + status = mgr.transfer_statuses[5] + self.assertEqual(status.received_kvs_per_pp, {}) + self.assertEqual(status.expected_kvs_per_pp, {}) + self.assertFalse(status.received_aux) + + +class TestGenerationBoundStagingControl(CustomTestCase): + @staticmethod + def _request(generation, *, legacy=False): + if legacy: + return [b"STAGING_REQ", b"7", b"0", b"1", b"peer", b"3"] + return [ + STAGING_REQ_WIRE_TAG, + generation.encode("ascii"), + b"7", + b"0", + b"1", + b"peer", + b"3", + ] + + def test_stale_and_legacy_staging_requests_cannot_allocate_reused_room(self): + sock = MagicMock() + receiver = SimpleNamespace( + transfer_generation=TRANSFER_GENERATION, + chunk_staging_infos=[], + _connect_to_bootstrap_server=MagicMock( + return_value=(sock, threading.Lock()) + ), + ) + allocator = SimpleNamespace( + assign=MagicMock(return_value=(3, 128, 0)), total_size=1 << 20 + ) + kwargs = dict( + staging_allocator=allocator, + kv_args=SimpleNamespace( + page_size=64, + kv_item_lens=[4096, 4096], + total_kv_head_num=4, + engine_rank=0, + ), + attn_tp_size=16, + prefill_attn_tp_size=1, + kv_buffer_tensors=None, + room_receivers={7: receiver}, + room_bootstrap={7: [{"pp_rank": 3}]}, + ) + + with patch.dict( + sys.modules, + { + "sglang.srt.disaggregation.common.staging_buffer": _fake_staging_buffer_module() + }, + ): + self.assertFalse( + handle_staging_req(self._request(OLD_TRANSFER_GENERATION), **kwargs) + ) + self.assertFalse( + handle_staging_req( + self._request(TRANSFER_GENERATION, legacy=True), **kwargs + ) + ) + + allocator.assign.assert_not_called() + receiver._connect_to_bootstrap_server.assert_not_called() + self.assertEqual(receiver.chunk_staging_infos, []) + + def test_staging_request_and_response_bind_exact_generation(self): + sock = MagicMock() + receiver = SimpleNamespace( + transfer_generation=TRANSFER_GENERATION, + chunk_staging_infos=[], + _connect_to_bootstrap_server=MagicMock( + return_value=(sock, threading.Lock()) + ), + ) + allocator = SimpleNamespace( + assign=MagicMock(return_value=(3, 128, 0)), total_size=1 << 20 + ) + with patch.dict( + sys.modules, + { + "sglang.srt.disaggregation.common.staging_buffer": _fake_staging_buffer_module() + }, + ): + self.assertTrue( + handle_staging_req( + self._request(TRANSFER_GENERATION), + allocator, + SimpleNamespace( + page_size=64, + kv_item_lens=[4096, 4096], + total_kv_head_num=4, + engine_rank=0, + ), + attn_tp_size=16, + prefill_attn_tp_size=1, + kv_buffer_tensors=None, + room_receivers={7: receiver}, + room_bootstrap={7: [{"pp_rank": 3}]}, + ) + ) + + expected_rsp = [ + STAGING_RSP_WIRE_TAG, + TRANSFER_GENERATION.encode("ascii"), + b"7", + b"0", + b"128", + b"0", + b"640", + b"peer", + ] + sock.send_multipart.assert_called_once_with(expected_rsp) + + tinfo = SimpleNamespace( + transfer_generation=TRANSFER_GENERATION, + staging=None, + ) + transfer_infos = {7: {"peer": tinfo}} + stale_rsp = list(expected_rsp) + stale_rsp[1] = OLD_TRANSFER_GENERATION.encode("ascii") + self.assertFalse(handle_staging_rsp(stale_rsp, transfer_infos)) + self.assertIsNone(tinfo.staging) + self.assertFalse( + handle_staging_rsp( + [b"STAGING_RSP", b"7", b"0", b"128", b"0", b"640", b"peer"], + transfer_infos, + ) + ) + self.assertIsNone(tinfo.staging) + + self.assertTrue(handle_staging_rsp(expected_rsp, transfer_infos)) + self.assertEqual(tinfo.staging.offsets, [128]) + self.assertEqual(tinfo.staging.rounds, [0]) + self.assertEqual(tinfo.staging.ends, [640]) + class TestNixlReceiverPoll(CustomTestCase): + def test_supplied_generation_is_shared_and_read_only_across_decode_ranks(self): + mgr = SimpleNamespace( + transfer_statuses={}, + addr_to_rooms_tracker=defaultdict(set), + update_status=MagicMock(), + ) + + first = NixlKVReceiver( + mgr, "prefill:8998", 11, transfer_generation=TRANSFER_GENERATION + ) + second = NixlKVReceiver( + mgr, "prefill:8998", 11, transfer_generation=TRANSFER_GENERATION + ) + + self.assertEqual(first.transfer_generation, TRANSFER_GENERATION) + self.assertEqual(second.transfer_generation, TRANSFER_GENERATION) + self.assertEqual( + mgr.transfer_statuses[11].transfer_generation, TRANSFER_GENERATION + ) + with self.assertRaises(AttributeError): + first.transfer_generation = OLD_TRANSFER_GENERATION + def _make_receiver(self, status=KVPoll.WaitingForInput): mgr = MagicMock() mgr.waiting_timeout = 5 @@ -698,6 +1421,7 @@ def _make_receiver(self, status=KVPoll.WaitingForInput): receiver.init_time = None receiver.conclude_state = None receiver.abort_notified = False + receiver._transfer_generation = TRANSFER_GENERATION receiver._connection_pool_entries = {} return receiver, mgr @@ -742,7 +1466,9 @@ def test_queued_completion_wins_over_waiting_timeout(self, mock_time): receiver, mgr = self._make_receiver(status=KVPoll.WaitingForInput) receiver.started_transfer = True receiver.init_time = 10.0 - mgr.transfer_statuses = {11: TransferStatus()} + mgr.transfer_statuses = { + 11: TransferStatus(transfer_generation=TRANSFER_GENERATION) + } mgr.check_transfer_done.return_value = True self.assertEqual(receiver.poll(), KVPoll.Success) @@ -757,7 +1483,7 @@ def test_transfer_done_returns_success_and_cleans_room_state(self, mock_time): receiver, mgr = self._make_receiver(status=KVPoll.WaitingForInput) receiver.started_transfer = True receiver.init_time = 10.0 - status = TransferStatus() + status = TransferStatus(transfer_generation=TRANSFER_GENERATION) status.received_aux = True status.num_pp_ranks_expected = 1 status.expected_kvs_per_pp[0] = 0 @@ -912,6 +1638,7 @@ def test_prefetch_staging_reqs_marks_room_when_no_peer_needs_staging(self): dst_aux_index=0, required_dst_info_num=1, dst_state_indices=[], + transfer_generation=TRANSFER_GENERATION, ) } } @@ -937,7 +1664,11 @@ def test_do_staging_transfer_requeues_when_allocation_not_ready(self): prefill_aux_index=None, state_indices=None, ) - req = SimpleNamespace(room=3, agent_name="decode_agent") + req = SimpleNamespace( + room=3, + agent_name="decode_agent", + completion_notification_prefix=f"nixlv2_{TRANSFER_GENERATION}_3", + ) queue = FakeQueue() with patch.dict( @@ -994,7 +1725,13 @@ def test_do_staging_transfer_raises_for_oversized_allocation(self): strategy, kv_chunk, kv_chunk.prefill_kv_indices, - SimpleNamespace(room=3, agent_name="decode_agent"), + SimpleNamespace( + room=3, + agent_name="decode_agent", + completion_notification_prefix=( + f"nixlv2_{TRANSFER_GENERATION}_3" + ), + ), SimpleNamespace(), FakeQueue(), ) @@ -1040,14 +1777,21 @@ def test_do_staging_transfer_builds_staging_notification(self): strategy, kv_chunk, kv_chunk.prefill_kv_indices, - SimpleNamespace(room=3, agent_name="decode_agent"), + SimpleNamespace( + room=3, + agent_name="decode_agent", + completion_notification_prefix=f"nixlv2_{TRANSFER_GENERATION}_3", + ), dst_info, FakeQueue(), ) self.assertEqual(handle, "handle") self.assertFalse(deferred) - self.assertEqual(calls[0][0][8], "3_stg_7_1_1_2_4_2_decode_agent") + self.assertEqual( + calls[0][0][8], + completion_notification(3, "stg_7_1_1_2_4_2_decode_agent"), + ) def test_send_kvcache_staged_uses_one_bulk_vram_write(self): mock_gather = MagicMock() @@ -1076,7 +1820,7 @@ def test_send_kvcache_staged_uses_one_bulk_vram_write(self): dst_tp_rank=0, dst_attn_tp_size=1, dst_kv_item_len=128, - notif="3_stg_0_1_1_0_0_2_decode_agent", + notif=completion_notification(3, "stg_0_1_1_0_0_2_decode_agent"), staging_buffer=FakeStagingBuffer(ptr=0x9000, size=1 << 20), ) @@ -1095,7 +1839,7 @@ def test_send_kvcache_staged_uses_one_bulk_vram_write(self): self.assertEqual(agent.initialize_xfer_calls[0][0], "WRITE") self.assertEqual( agent.initialize_xfer_calls[0][-1], - b"3_stg_0_1_1_0_0_2_decode_agent", + completion_notification(3, "stg_0_1_1_0_0_2_decode_agent").encode("ascii"), ) def test_send_kvcache_staged_falls_back_when_prefill_buffer_too_small(self): @@ -1141,6 +1885,44 @@ def prep_xfer_dlist(self, peer_name, array, mem_kind): return f"handle_{len(self.calls)}" +class TestNixlTorchTransferCompactDlist(CustomTestCase): + def test_equal_tp_dlist_uses_one_strided_region_per_allocation(self): + mgr = object.__new__(NixlKVManager) + mgr._use_torch_transfer = True + mgr.agent = DlistCaptureAgent() + + handle = mgr._prep_equal_tp_dlist( + "decode-agent", + kv_ptrs=[0x10000, 0x20000], + kv_item_lens=[64, 128], + kv_data_lens=[640, 1280], + gpu_id=3, + ) + + self.assertEqual(handle, "handle_1") + peer_name, descriptors, memory_type = mgr.agent.calls[0] + self.assertEqual(peer_name, "decode-agent") + self.assertEqual(memory_type, "VRAM") + np.testing.assert_array_equal( + descriptors, + np.asarray( + [ + [0x10000, 64, 3, 64, 10], + [0x20000, 128, 3, 128, 10], + ], + dtype=np.uint64, + ), + ) + np.testing.assert_array_equal( + repeat_indices_over_layers( + np.asarray([1, 9], dtype=np.int32), + num_layers=2, + layer_length=10, + ), + np.asarray([1, 9, 11, 19], dtype=np.int32), + ) + + class TestNixlHeteroTpReplicatedKV(CustomTestCase): """Regression guard for #31295. 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 752787665e3b..c88840de147c 100644 --- a/test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py +++ b/test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py @@ -7,78 +7,163 @@ after its DONE barrier -- never from the bootstrap thread for an active room. """ +import threading import unittest from unittest.mock import MagicMock from sglang.srt.disaggregation.base.conn import KVPoll from sglang.srt.disaggregation.common.conn import CommonKVManager -from sglang.srt.disaggregation.nixl.conn import NixlKVManager +from sglang.srt.disaggregation.nixl.conn import GUARD, 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") +GEN_A = "11" * 16 +GEN_B = "22" * 16 + 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._abort_state_lock = threading.Lock() mgr._deferred_ack_targets = {} + mgr._active_transfer_generations = {} + mgr._abort_ack_receipts = set() mgr._staging_outstanding = {} mgr.request_status = {} mgr._sent = [] # Capture acks instead of opening a socket. - mgr._send_abort_ack = lambda ip, port, room: mgr._sent.append((ip, port, room)) + mgr._send_abort_ack = lambda ip, port, room, generation: ( + mgr._sent.append((ip, port, room, generation)) or True + ) return mgr +def _activate_failed(mgr, room, generation=GEN_A): + mgr.register_transfer_generation(room, generation) + mgr.request_status[room] = KVPoll.Failed + return mgr.register_deferred_ack_target(room, generation, "10.0.0.1", 5000) + + 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) + self.assertTrue(_activate_failed(mgr, 7)) mgr._staging_outstanding[7] = 1 - mgr._maybe_ack_drained_abort(7) + mgr._maybe_ack_drained_abort(7, GEN_A) 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)]) + mgr._maybe_ack_drained_abort(7, GEN_A) + self.assertEqual(mgr._sent, [("10.0.0.1", 5000, 7, GEN_A)]) + + def test_one_skipped_writer_cannot_erase_an_active_sibling(self): + mgr = _prefill_mgr() + self.assertTrue(_activate_failed(mgr, 70)) + self.assertEqual(mgr._count_transfer_writer(70), 1) + self.assertEqual(mgr._count_transfer_writer(70), 2) + + self.assertEqual(mgr._uncount_transfer_writer(70), 1) + mgr._maybe_ack_drained_abort(70, GEN_A) + self.assertEqual(mgr._sent, []) + self.assertFalse(mgr.mark_transfer_quiescent(70, GEN_A)) + self.assertEqual(mgr.active_transfer_generation(70), GEN_A) + + self.assertEqual(mgr._uncount_transfer_writer(70), 0) + self.assertTrue(mgr.mark_transfer_quiescent(70, GEN_A)) + self.assertEqual(mgr._sent, [("10.0.0.1", 5000, 70, GEN_A)]) + self.assertIsNone(mgr.active_transfer_generation(70)) 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) + mgr.register_transfer_generation(8, GEN_A) + mgr.request_status[8] = KVPoll.Failed + mgr.register_deferred_ack_target(8, GEN_A, "10.0.0.2", 5001) + mgr._maybe_ack_drained_abort(8, GEN_A) + mgr._maybe_ack_drained_abort(8, GEN_A) self.assertEqual(len(mgr._sent), 1) - self.assertNotIn(8, mgr._deferred_ack_targets) + self.assertNotIn((8, GEN_A), mgr._deferred_ack_targets) def test_unregistered_room_is_noop(self): mgr = _prefill_mgr() - mgr._maybe_ack_drained_abort(999) + mgr._maybe_ack_drained_abort(999, GEN_A) self.assertEqual(mgr._sent, []) + def test_normal_send_failure_retains_target_and_receipt_for_retry(self): + mgr = _prefill_mgr() + self.assertTrue(_activate_failed(mgr, 9)) + mgr._send_abort_ack = MagicMock(side_effect=[False, True]) + + mgr._maybe_ack_drained_abort(9, GEN_A) + self.assertIn((9, GEN_A), mgr._deferred_ack_targets) + self.assertIn((9, GEN_A), mgr._abort_ack_receipts) + mgr._maybe_ack_drained_abort(9, GEN_A) + self.assertNotIn((9, GEN_A), mgr._deferred_ack_targets) + self.assertEqual(mgr._send_abort_ack.call_count, 2) + + def test_base_exception_after_acceptance_retains_durable_retry_owner(self): + mgr = _prefill_mgr() + self.assertTrue(_activate_failed(mgr, 10)) + accepted = [] + + def accepted_then_interrupted(ip, port, room, generation): + accepted.append((ip, port, room, generation)) + raise KeyboardInterrupt("lost return") + + mgr._send_abort_ack = accepted_then_interrupted + with self.assertRaises(KeyboardInterrupt): + mgr._maybe_ack_drained_abort(10, GEN_A) + self.assertIn((10, GEN_A), mgr._deferred_ack_targets) + self.assertIn((10, GEN_A), mgr._abort_ack_receipts) + + mgr._send_abort_ack = lambda ip, port, room, generation: ( + accepted.append((ip, port, room, generation)) or True + ) + mgr._maybe_ack_drained_abort(10, GEN_A) + self.assertNotIn((10, GEN_A), mgr._deferred_ack_targets) + self.assertEqual(len(accepted), 2) + + def test_receipt_tombstones_are_not_evicted(self): + mgr = _prefill_mgr() + for room in range(2048): + mgr._remember_abort_receipt_unlocked((room, GEN_A)) + self.assertEqual(len(mgr._abort_ack_receipts), 2048) + self.assertIn((0, GEN_A), mgr._abort_ack_receipts) + 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) + def test_ack_target_is_immutable_for_one_generation(self): + mgr = _prefill_mgr() + mgr.register_transfer_generation(12, GEN_A) + self.assertTrue(mgr.register_deferred_ack_target(12, GEN_A, "10.0.0.1", 5000)) + self.assertFalse(mgr.register_deferred_ack_target(12, GEN_A, "10.0.0.2", 5001)) + target = mgr._deferred_ack_targets[(12, GEN_A)] + self.assertEqual((target.decode_ip, target.decode_port), ("10.0.0.1", 5000)) + 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): + def _abort_msg(room=11, ip="10.0.0.3", port=6000, generation=GEN_A): return [ b"ABORT", str(room).encode("ascii"), ip.encode("ascii"), str(port).encode("ascii"), + generation.encode("ascii"), ] def _mgr(self, enabled=True, room=11, status=KVPoll.WaitingForInput): mgr = _prefill_mgr(NixlKVManager, enabled=enabled) + mgr.register_transfer_generation(room, GEN_A) if status is not None: mgr.request_status[room] = status mgr.record_failure = MagicMock() @@ -94,7 +179,8 @@ def test_in_flight_room_registers_target_and_does_not_ack_yet(self): 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)) + target = mgr._deferred_ack_targets[(11, GEN_A)] + self.assertEqual((target.decode_ip, target.decode_port), ("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) @@ -105,7 +191,7 @@ def test_quiescent_active_room_acks_without_waiting_for_a_worker_visit(self): 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._sent, [("10.0.0.3", 6000, 11, GEN_A)]) self.assertEqual(mgr._deferred_ack_targets, {}) def test_worker_skip_before_registration_still_acks(self): @@ -120,21 +206,22 @@ 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.mark_transfer_quiescent(room, GEN_A) 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._sent, [("10.0.0.3", 6000, 11, GEN_A)]) 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.mark_transfer_quiescent(11, GEN_A) 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._sent, [("10.0.0.3", 6000, 11, GEN_A)]) self.assertEqual(mgr._deferred_ack_targets, {}) def test_cleared_room_with_outstanding_chunk_does_not_ack(self): @@ -146,17 +233,26 @@ def test_cleared_room_with_outstanding_chunk_does_not_ack(self): self.assertTrue(mgr._handle_abort_notification(self._abort_msg())) self.assertEqual(mgr._sent, []) - self.assertEqual(mgr._deferred_ack_targets, {}) + self.assertIn((11, GEN_A), mgr._deferred_ack_targets) - def test_feature_off_registers_nothing_and_acks_nothing(self): + def test_explicit_feature_false_cannot_disable_safe_ack_protocol(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._sent, [("10.0.0.3", 6000, 11, GEN_A)]) self.assertEqual(mgr.request_status[11], KVPoll.Failed) + def test_old_generation_abort_cannot_mutate_reused_room(self): + mgr = self._mgr() + self.assertTrue( + mgr._handle_abort_notification(self._abort_msg(generation=GEN_B)) + ) + self.assertEqual(mgr.request_status[11], KVPoll.WaitingForInput) + self.assertEqual(mgr.active_transfer_generation(11), GEN_A) + self.assertEqual(mgr._deferred_ack_targets, {}) + self.assertEqual(mgr._sent, []) + def test_legacy_two_frame_abort_is_tolerated(self): # Older peers send [ABORT, room] with no return address. mgr = self._mgr() @@ -171,17 +267,139 @@ def test_non_abort_message_is_not_claimed(self): 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 = _prefill_mgr() mgr._deferred_abort_ack_tracker = {} - mgr.register_deferred_abort_room(21) + mgr.register_deferred_abort_room(21, GEN_A, frozenset({0, 1})) + + for rank in (0, 1, 1): + mgr.note_abort_ack(21, GEN_A, rank) - 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())) + self.assertTrue(mgr.is_abort_release_safe(21, GEN_A)) + + def test_late_metadata_does_not_poison_new_generation(self): + mgr = _prefill_mgr(NixlKVManager) + mgr._bootstrap_errors = {} + mgr.exceptions = {} + mgr.transfer_infos = {21: {}} + mgr.record_failure = MagicMock() + mgr.register_transfer_generation(21, GEN_B) + late_payload = [ + b"21", + b"endpoint", + b"1234", + b"agent", + b"", + b"0", + b"1", + b"", + b"", + b"0", + GEN_A.encode("ascii"), + ] + + mgr._handle_bootstrap_message([GUARD, *late_payload]) + + self.assertEqual(mgr.active_transfer_generation(21), GEN_B) + self.assertEqual(mgr.transfer_infos[21], {}) + self.assertNotIn(21, mgr.exceptions) + mgr.record_failure.assert_not_called() + + def test_malformed_or_generationless_metadata_cannot_poison_reused_room(self): + for generation_frame in (b"not-a-generation", None): + mgr = _prefill_mgr(NixlKVManager) + mgr._bootstrap_errors = {} + mgr.exceptions = {} + mgr.transfer_infos = {23: {}} + mgr.record_failure = MagicMock() + mgr.register_transfer_generation(23, GEN_B) + payload = [ + b"23", + b"endpoint", + b"1234", + b"agent", + b"", + b"0", + b"1", + b"", + b"", + b"0", + ] + if generation_frame is not None: + payload.append(generation_frame) + + mgr._handle_bootstrap_message([GUARD, *payload]) + + self.assertEqual(mgr.active_transfer_generation(23), GEN_B) + self.assertEqual(mgr.transfer_infos[23], {}) + self.assertNotIn(23, mgr.exceptions) + self.assertNotIn("agent", mgr._bootstrap_errors) + mgr.record_failure.assert_not_called() + + def test_retired_metadata_cannot_resurrect_room_before_reuse(self): + mgr = _prefill_mgr(NixlKVManager) + mgr._bootstrap_errors = {} + mgr.exceptions = {} + mgr.transfer_infos = {} + mgr.record_failure = MagicMock() + mgr.register_transfer_generation(22, GEN_A) + mgr.mark_transfer_quiescent(22, GEN_A) + late_payload = [ + b"22", + b"endpoint", + b"1234", + b"agent", + b"", + b"0", + b"1", + b"", + b"", + b"0", + GEN_A.encode("ascii"), + ] - self.assertFalse(mgr.is_abort_release_safe(21, required_acks=3)) - self.assertTrue(mgr.is_abort_release_safe(21, required_acks=2)) + mgr._handle_bootstrap_message([GUARD, *late_payload]) + + self.assertIsNone(mgr.active_transfer_generation(22)) + self.assertNotIn(22, mgr.transfer_infos) + self.assertIn((22, GEN_A), mgr._abort_ack_receipts) + self.assertNotIn(22, mgr.exceptions) + mgr.record_failure.assert_not_called() + + def test_retirement_and_duplicate_registration_are_atomic(self): + mgr = _prefill_mgr(NixlKVManager) + rooms = tuple(range(1000, 1256)) + for room in rooms: + self.assertTrue(mgr.register_transfer_generation(room, GEN_A)) + + operation_start = threading.Barrier(2) + operation_done = threading.Barrier(2) + replay_results = [] + + def retire_generations(): + for room in rooms: + operation_start.wait() + mgr.mark_transfer_quiescent(room, GEN_A) + operation_done.wait() + + def replay_metadata_registration(): + for room in rooms: + operation_start.wait() + replay_results.append(mgr.register_transfer_generation(room, GEN_A)) + operation_done.wait() + + retire_thread = threading.Thread(target=retire_generations) + replay_thread = threading.Thread(target=replay_metadata_registration) + retire_thread.start() + replay_thread.start() + retire_thread.join(timeout=5) + replay_thread.join(timeout=5) + + self.assertFalse(retire_thread.is_alive()) + self.assertFalse(replay_thread.is_alive()) + self.assertEqual(len(replay_results), len(rooms)) + for room in rooms: + self.assertIsNone(mgr.active_transfer_generation(room)) + self.assertIn((room, GEN_A), mgr._abort_ack_receipts) 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..54cdd35ee2e7 100644 --- a/test/registered/unit/disaggregation/test_nixl_sender_failure_cleanup.py +++ b/test/registered/unit/disaggregation/test_nixl_sender_failure_cleanup.py @@ -12,6 +12,7 @@ class TestNixlSenderFailureCleanup(unittest.TestCase): def test_failure_exception_cleans_room_state_before_raising(self): room = 7 + generation = "0123456789abcdef0123456789abcdef" expected_exc = RuntimeError("transfer failed") sender = NixlKVSender.__new__(NixlKVSender) sender.bootstrap_room = room @@ -20,7 +21,10 @@ def test_failure_exception_cleans_room_state_before_raising(self): sender._send_error = None staging_ctx = SimpleNamespace( prefetched_rooms={room, 8}, - prefetch_requested={(room, 0, "session-a"), (8, 0, "session-b")}, + prefetch_requested={ + (room, generation, 0, "session-a"), + (8, generation, 0, "session-b"), + }, ) sender.kv_mgr = SimpleNamespace( enable_staging=True, @@ -45,9 +49,11 @@ def test_failure_exception_cleans_room_state_before_raising(self): self.assertNotIn(room, sender.kv_mgr.exceptions) self.assertNotIn(room, sender.kv_mgr.failure_records) self.assertNotIn(room, staging_ctx.prefetched_rooms) - self.assertNotIn((room, 0, "session-a"), staging_ctx.prefetch_requested) + self.assertNotIn( + (room, generation, 0, "session-a"), staging_ctx.prefetch_requested + ) self.assertIn(8, staging_ctx.prefetched_rooms) - self.assertIn((8, 0, "session-b"), staging_ctx.prefetch_requested) + self.assertIn((8, generation, 0, "session-b"), staging_ctx.prefetch_requested) if __name__ == "__main__": diff --git a/test/registered/unit/disaggregation/test_nixl_torch_transfer.py b/test/registered/unit/disaggregation/test_nixl_torch_transfer.py new file mode 100644 index 000000000000..3c63b1d7bb80 --- /dev/null +++ b/test/registered/unit/disaggregation/test_nixl_torch_transfer.py @@ -0,0 +1,2257 @@ +"""CPU tests for the experimental PyTorch transfer compatibility boundary.""" + +# Thread targets intentionally retain BaseException for process-control fault +# injection; nested contexts keep the interruption point visually explicit. +# ruff: noqa: BLE001, SIM117 + +from __future__ import annotations + +import ast +import threading +import time +from enum import Enum, auto +from pathlib import Path +from types import SimpleNamespace +from typing import ClassVar +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +from sglang.srt.disaggregation.nixl.torch_transfer import ( + TorchTransferAgent, + _load_transfer_api, + _register_nixl_backend, + _WorkHandle, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class _WorkState(Enum): + PENDING = auto() + RUNNING = auto() + COMPLETED = auto() + SUCCEEDED = COMPLETED + FAILED = auto() + CANCELLED = auto() + + +class _TransferOp(Enum): + READ = "read" + WRITE = "write" + + +class _Resource: + def __init__(self) -> None: + self.closed = False + self.close_calls = 0 + self.close_errors: list[BaseException] = [] + + def close(self) -> None: + self.close_calls += 1 + if self.close_errors: + raise self.close_errors.pop(0) + self.closed = True + + +class _RawSpan: + instances: ClassVar[list[_RawSpan]] = [] + + def __init__(self, **kwargs) -> None: + self.__dict__.update(kwargs) + self.__class__.instances.append(self) + + +class _Registration(_Resource): + def __init__(self, span: _RawSpan, name: str) -> None: + super().__init__() + self.span = span + self.name = name + self.region_calls = [] + + def region(self, offset: int, nbytes: int, **kwargs): + region = ("local", self.name, offset, nbytes, kwargs) + self.region_calls.append(region) + return region + + +class _Peer(_Resource): + def __init__(self, name: str = "remote-endpoint") -> None: + super().__init__() + self.name = name + self.endpoint_id = "remote-id" + self.incarnation = "remote-incarnation" + self.region_calls = [] + + def unsafe_region(self, address: int, nbytes: int, **kwargs): + region = ("remote", address, nbytes, kwargs) + self.region_calls.append(region) + return region + + +class _Plan(_Resource): + def __init__(self, endpoint, kwargs: dict) -> None: + super().__init__() + self.endpoint = endpoint + self.kwargs = kwargs + self.select_indices_calls = [] + self.submit_and_poll_calls = [] + + def select_indices( + self, + operation, + *, + local_indices, + remote_indices, + ): + self.select_indices_calls.append((operation, local_indices, remote_indices)) + raise AssertionError("one-shot SGLang transfers must not create IndexSelection") + + def submit_indices_and_poll( + self, + operation, + *, + local_indices, + remote_indices, + **kwargs, + ): + local_snapshot = np.asarray(local_indices, dtype=np.int32).copy() + remote_snapshot = np.asarray(remote_indices, dtype=np.int32).copy() + call = { + "operation": operation, + "local_indices": tuple(local_snapshot), + "remote_indices": tuple(remote_snapshot), + **kwargs, + } + self.submit_and_poll_calls.append(call) + submission_kwargs = { + key: value + for key, value in kwargs.items() + if key not in {"max_polls", "timeout_ns"} + } + work = self.endpoint._submit( + operation, + self, + { + "local_indices": local_snapshot, + "remote_indices": remote_snapshot, + **submission_kwargs, + }, + ) + endpoint = self.endpoint + state = ( + endpoint.next_fused_state + if endpoint.capabilities.fused_prevalidated_submit_poll + else _WorkState.PENDING + ) + work.done = state in ( + _WorkState.COMPLETED, + _WorkState.FAILED, + _WorkState.CANCELLED, + ) + work.state = state + work.error = endpoint.next_fused_error + endpoint._raise_postcommit_error() + return work, state + + +class _Selection: + def __init__(self, plan, operation, local_indices, remote_indices) -> None: + self.plan = plan + self.operation = operation + self.local_indices = np.asarray(local_indices, dtype=np.int32).copy() + self.remote_indices = np.asarray(remote_indices, dtype=np.int32).copy() + self.submit_and_poll_calls = [] + + def submit_and_poll(self, **kwargs): + call = dict(kwargs) + self.submit_and_poll_calls.append(call) + self.plan.submit_and_poll_calls.append((self, call)) + submission_kwargs = { + key: value + for key, value in kwargs.items() + if key not in {"max_polls", "timeout_ns"} + } + work = self.plan.endpoint._submit( + self.operation, + self.plan, + { + "local_indices": self.local_indices, + "remote_indices": self.remote_indices, + **submission_kwargs, + }, + ) + endpoint = self.plan.endpoint + state = ( + endpoint.next_fused_state + if endpoint.capabilities.fused_prevalidated_submit_poll + else _WorkState.PENDING + ) + work.done = state in ( + _WorkState.COMPLETED, + _WorkState.FAILED, + _WorkState.CANCELLED, + ) + work.state = state + work.error = endpoint.next_fused_error + endpoint._raise_postcommit_error() + return work, state + + +class _Work(_Resource): + def __init__(self, recovery_token=None) -> None: + super().__init__() + self.recovery_token = recovery_token + self.done = False + self._state = _WorkState.PENDING + self.error = None + self.cancel_completes = True + self.test_error: Exception | None = None + self.test_calls = 0 + self.state_hook = None + + @property + def state(self) -> _WorkState: + self.test_calls += 1 + if self.state_hook is not None: + self.state_hook() + if self.test_error is not None: + raise self.test_error + return self._state + + @state.setter + def state(self, value: _WorkState) -> None: + self._state = value + + def test(self) -> bool: + self.test_calls += 1 + if self.test_error is not None: + raise self.test_error + return self.done + + def cancel(self) -> bool: + if self.cancel_completes: + self.done = True + self.state = _WorkState.CANCELLED + return self.cancel_completes + + def close(self) -> None: + super().close() + if self.closed: + self.recovery_token = None + + +class _ProgressMode(Enum): + MANUAL = auto() + BACKGROUND = auto() + + +class _ThreadMode(Enum): + CALLER_SERIALIZED = auto() + SERIALIZED = auto() + MULTIPLE = auto() + + +class _Endpoint(_Resource): + instances: ClassVar[list[_Endpoint]] = [] + synchronous_notification_send = True + constructor_cut: tuple[BaseException, bool, list[BaseException]] | None = None + missing_capability: str | None = None + + def __init__( + self, + name: str, + *, + backend: str, + endpoint_id: str, + progress_mode: _ProgressMode, + thread_mode: _ThreadMode, + options: dict, + ) -> None: + super().__init__() + self.name = name + self.backend = backend + self.endpoint_id = endpoint_id + self.progress_mode = progress_mode + self.thread_mode = thread_mode + self.options = options + self.capabilities = SimpleNamespace( + notifications=True, + attached_notifications=True, + standalone_notifications=True, + raw_spans=True, + strided_regions=True, + write=True, + fused_prevalidated_submit_poll=True, + synchronous_notification_send=self.__class__.synchronous_notification_send, + ) + self.registrations = [] + self.peers = [] + self.plans = [] + self.bound_transfers = [] + self.transfer_slots = [] + self.submissions = [] + self.notifications = [] + self.synchronous_notifications = [] + self.incoming_notifications = [] + self.progress_calls = 0 + self.next_fused_result = False + self.next_fused_state = _WorkState.PENDING + self.next_fused_error = None + self.poison_index_submit = False + self.submit_attempts = 0 + self.submit_error: BaseException | None = None + self.postcommit_error: BaseException | None = None + self.postcommit_hook = None + self.open_works_errors: list[BaseException] = [] + self.next_work_done = False + self.next_work_state = _WorkState.PENDING + self._open_works = [] + self.open_works_reads = 0 + self.construction_snapshot_reads = 0 + self.poison_construction_snapshots = False + missing_capability = self.__class__.missing_capability + if missing_capability is not None: + setattr(self.capabilities, missing_capability, False) + self.__class__.instances.append(self) + constructor_cut = self.__class__.constructor_cut + if constructor_cut is not None: + self.__class__.constructor_cut = None + error, mismatch, close_errors = constructor_cut + self.close_errors.extend(close_errors) + if mismatch: + self.endpoint_id = f"foreign-{endpoint_id}" + error.recovery_endpoint = self + raise error + + @property + def open_registrations(self): + self._record_construction_snapshot() + return tuple(item for item in self.registrations if not item.closed) + + @property + def open_peers(self): + self._record_construction_snapshot() + return tuple(item for item in self.peers if not item.closed) + + @property + def open_plans(self): + self._record_construction_snapshot() + return tuple(item for item in self.plans if not item.closed) + + @property + def open_bound_transfers(self): + self._record_construction_snapshot() + return tuple(item for item in self.bound_transfers if not item.closed) + + @property + def open_transfer_slots(self): + self._record_construction_snapshot() + return tuple(item for item in self.transfer_slots if not item.closed) + + def _record_construction_snapshot(self): + self.construction_snapshot_reads += 1 + if self.poison_construction_snapshots: + raise AssertionError("clean path read a Core construction snapshot") + + @property + def open_works(self): + self.open_works_reads += 1 + if self.open_works_errors: + raise self.open_works_errors.pop(0) + return tuple(work for work in self._open_works if not work.closed) + + def register(self, span, *, name): + registration = _Registration(span, name) + self.registrations.append(registration) + return registration + + def export_metadata(self, registrations): + assert registrations == list(self.open_registrations) + return b"endpoint-metadata" + + def import_peer(self, metadata, **kwargs): + assert metadata == b"peer-metadata" + if kwargs: + assert kwargs == {"expected_name": "decode-agent"} + peer = _Peer() + self.peers.append(peer) + return peer + + def prepare(self, **kwargs): + plan = _Plan(self, kwargs) + self.plans.append(plan) + return plan + + def _submit(self, operation, plan, kwargs): + self.submit_attempts += 1 + if self.submit_error is not None: + raise self.submit_error + work = _Work(kwargs.get("recovery_token")) + work.done = self.next_work_done + work.state = self.next_work_state + self.submissions.append((operation, plan, kwargs, work)) + self._open_works.append(work) + return work + + def _raise_postcommit_error(self): + if self.postcommit_hook is not None: + self.postcommit_hook() + error, self.postcommit_error = self.postcommit_error, None + if error is not None: + raise error + + def write(self, plan, **kwargs): + work = self._submit("WRITE", plan, kwargs) + self._raise_postcommit_error() + return work + + def read(self, plan, **kwargs): + work = self._submit("READ", plan, kwargs) + self._raise_postcommit_error() + return work + + def progress(self): + self.progress_calls += 1 + + def send_notification(self, peer, payload, *, recovery_token=None): + work = _Work(recovery_token) + self.notifications.append((peer, payload, recovery_token, work)) + self._open_works.append(work) + self._raise_postcommit_error() + return work + + def send_notification_sync(self, peer, payload): + self.synchronous_notifications.append((peer, payload)) + + def poll_notifications(self): + result = self.incoming_notifications + self.incoming_notifications = [] + return result + + +def _make_api(): + return SimpleNamespace( + BACKEND_FACTORY_API_VERSION=2, + BackendFactoryV2=object, + Endpoint=_Endpoint, + ProgressMode=_ProgressMode, + RawSpan=_RawSpan, + ThreadMode=_ThreadMode, + TransferOp=_TransferOp, + TransferPlan=_Plan, + IndexSelection=_Selection, + WorkState=_WorkState, + ) + + +def _reset_fake_api() -> None: + _Endpoint.instances.clear() + _Endpoint.constructor_cut = None + _Endpoint.missing_capability = None + _Endpoint.synchronous_notification_send = True + _RawSpan.instances.clear() + + +def _make_uncreated_agent(): + _reset_fake_api() + with patch( + "sglang.srt.disaggregation.nixl.torch_transfer._load_transfer_api", + return_value=_make_api(), + ): + return TorchTransferAgent( + "local-endpoint", + owner=object(), + background_progress=False, + endpoint_backend="fake", + ) + + +def _make_agent( + *, + background_progress: bool = False, + synchronous_notification_send: bool = True, +): + _reset_fake_api() + _Endpoint.synchronous_notification_send = synchronous_notification_send + api = _make_api() + owner = object() + with ( + patch( + "sglang.srt.disaggregation.nixl.torch_transfer._load_transfer_api", + return_value=api, + ), + patch( + "sglang.srt.disaggregation.nixl.torch_transfer._register_nixl_backend" + ) as register_nixl, + ): + agent = TorchTransferAgent( + "local-endpoint", + owner=owner, + background_progress=background_progress, + num_threads=8 if background_progress else 0, + ) + agent.create_backend("UCX", {"num_threads": "8"}) + register_nixl.assert_called_once_with() + return agent, _Endpoint.instances[-1], owner + + +def _register_and_add_peer(agent): + registrations = agent.register_memory([(0x1000, 0x400, 0, "")], "VRAM") + agent.add_remote_agent(b"peer-metadata", peer_name="decode-agent") + return registrations + + +def test_raw_registration_keeps_owner_and_exports_metadata(): + agent, endpoint, owner = _make_agent(background_progress=True) + + registrations = agent.register_memory( + [(0x1000, 0x100, 2, ""), (0x2000, 0x80, 0, "")], "VRAM" + ) + + assert endpoint.progress_mode is _ProgressMode.BACKGROUND + assert endpoint.thread_mode is _ThreadMode.MULTIPLE + assert endpoint.options == { + "backends": ["UCX"], + "backend_init_params": {"UCX": {"num_threads": "8"}}, + "transfer_backends": ["UCX"], + "num_threads": 8, + } + assert len(registrations) == 2 + assert _RawSpan.instances[0].device == "cuda:2" + assert _RawSpan.instances[0].owner is owner + assert agent.get_agent_metadata() == b"endpoint-metadata" + + +def test_prepared_write_reuses_indexed_plan_and_attaches_notification(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.prep_xfer_dlist("", [(0x1000, 64, 0), (0x1040, 64, 0)], "VRAM") + remote = agent.prep_xfer_dlist( + "decode-agent", [(0x5000, 64, 0), (0x5040, 64, 0)], "VRAM" + ) + + source_indices = np.asarray([0, 1], dtype=np.int32) + destination_indices = np.asarray([1, 0], dtype=np.int32) + notification = bytearray(b"7_kv_0_1_0") + first = agent.make_prepped_xfer( + "WRITE", + local, + source_indices, + remote, + destination_indices, + notification, + ) + second = agent.make_prepped_xfer("WRITE", local, [1], remote, [0]) + + assert len(endpoint.plans) == 1 + assert endpoint.plans[0].kwargs["indexed"] is True + assert endpoint.submissions == [] + assert endpoint.submit_attempts == 0 + assert first.work is None + assert second.work is None + assert endpoint.plans[0].select_indices_calls == [] + assert first.local_indices is source_indices + assert first.remote_indices is destination_indices + assert source_indices.flags.writeable is False + assert destination_indices.flags.writeable is False + + with pytest.raises(ValueError, match="read-only"): + source_indices[:] = [1, 1] + with pytest.raises(ValueError, match="read-only"): + destination_indices[:] = [0, 0] + notification[:] = b"changed_tag" + + assert agent.transfer(first) == "PROC" + first_submission = endpoint.submissions[0][2] + assert tuple(first_submission["local_indices"]) == (0, 1) + assert tuple(first_submission["remote_indices"]) == (1, 0) + assert first_submission["local_indices"].dtype == np.int32 + assert first_submission["remote_indices"].dtype == np.int32 + assert first_submission["notification"] == b"7_kv_0_1_0" + assert endpoint.plans[0].submit_and_poll_calls[0] == { + "operation": _TransferOp.WRITE, + "local_indices": (0, 1), + "remote_indices": (1, 0), + "notification": b"7_kv_0_1_0", + "max_polls": 0, + "timeout_ns": None, + "recovery_token": first, + } + assert first.local_indices is None + assert first.remote_indices is None + assert endpoint.plans[0].select_indices_calls == [] + assert endpoint.open_works_reads == 0 + assert first.work.test_calls == 0 + first.work.done = True + first.work.state = _WorkState.SUCCEEDED + assert agent.check_xfer_state(first) == "DONE" + assert first.work.test_calls == 1 + assert first.work.closed is True + + assert agent.transfer(second) == "PROC" + assert second.work.test_calls == 0 + assert endpoint.submit_attempts == 2 + assert agent.transfer(second) == "PROC" + assert second.work.test_calls == 1 + assert endpoint.submit_attempts == 2 + second.work.done = True + second.work.state = _WorkState.SUCCEEDED + assert agent.check_xfer_state(second) == "DONE" + assert second.work.test_calls == 2 + agent.release_dlist_handle(local) + assert endpoint.plans[0].closed is True + + +@pytest.mark.parametrize( + ("state", "error", "expected"), + [ + (_WorkState.SUCCEEDED, None, "DONE"), + (_WorkState.FAILED, RuntimeError("fused failure"), "ERR"), + ], +) +def test_indexed_fused_submit_consumes_immediate_terminal_without_state_reread( + state, error, expected +): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + endpoint.next_fused_result = True + endpoint.next_fused_state = state + endpoint.next_fused_error = error + endpoint.poison_index_submit = True + local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") + remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") + + handle = agent.make_prepped_xfer("WRITE", local, [0], remote, [0], b"fused") + assert endpoint.submissions == [] + assert handle.work is None + assert agent.transfer(handle) == expected + assert endpoint.plans[-1].select_indices_calls == [] + assert endpoint.plans[-1].submit_and_poll_calls == [ + { + "operation": _TransferOp.WRITE, + "local_indices": (0,), + "remote_indices": (0,), + "notification": b"fused", + "max_polls": 0, + "timeout_ns": None, + "recovery_token": handle, + } + ] + assert handle.local_indices is None + assert handle.remote_indices is None + assert handle.work.test_calls == 0 + if expected == "ERR": + assert handle.work.closed is True + assert id(handle) not in agent._works + assert agent.check_xfer_state(handle) == expected + assert handle.work.test_calls == 0 + assert handle.work.closed is True + + +def test_work_state_is_one_authoritative_read_then_cached(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") + assert handle.work is None + assert endpoint.submissions == [] + endpoint.next_work_done = True + endpoint.next_work_state = _WorkState.SUCCEEDED + + assert agent.transfer(handle) == "DONE" + handle.work.test = MagicMock(side_effect=AssertionError("test() was queried")) + assert handle.work.test_calls == 1 + assert handle.work.test.call_count == 0 + assert agent.check_xfer_state(handle) == "DONE" + assert handle.work.test_calls == 1 + assert handle.work.closed is True + + +def test_indexed_submit_fallback_is_owned_by_core_api(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + endpoint.capabilities.fused_prevalidated_submit_poll = False + local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") + remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") + + handle = agent.make_prepped_xfer("WRITE", local, [0], remote, [0], b"fallback") + assert endpoint.submissions == [] + assert endpoint.plans[-1].submit_and_poll_calls == [] + assert handle.first_state is None + assert agent.transfer(handle) == "PROC" + assert endpoint.plans[-1].select_indices_calls == [] + assert endpoint.plans[-1].submit_and_poll_calls == [ + { + "operation": _TransferOp.WRITE, + "local_indices": (0,), + "remote_indices": (0,), + "notification": b"fallback", + "max_polls": 0, + "timeout_ns": None, + "recovery_token": handle, + } + ] + assert handle.first_state is None + assert handle.work.test_calls == 0 + assert agent.transfer(handle) == "PROC" + assert handle.work.test_calls == 1 + assert endpoint.submit_attempts == 1 + + +def test_plan_close_failure_quarantines_cache_and_peer_until_retry(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") + remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") + handle = agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) + assert agent.transfer(handle) == "PROC" + handle.work.done = True + handle.work.state = _WorkState.SUCCEEDED + assert agent.check_xfer_state(handle) == "DONE" + plan = endpoint.plans[-1] + plan.close_errors.append(RuntimeError("partially committed close")) + + with pytest.raises(RuntimeError, match="partially committed"): + agent.remove_remote_agent("decode-agent") + + assert not agent._plan_cache + assert agent._closing_plans + assert "decode-agent" in agent._retiring_peers + with pytest.raises(RuntimeError, match="teardown is in progress"): + agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) + + agent.remove_remote_agent("decode-agent") + assert not agent._closing_plans + assert "decode-agent" not in agent._retiring_peers + assert plan.closed is True + + +def test_compact_strided_catalog_preserves_flattened_indices(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.prep_xfer_dlist("", [(0x1000, 32, 0, 64, 8)], "VRAM") + remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0, 64, 8)], "VRAM") + + handle = agent.make_prepped_xfer("WRITE", local, [7, 1], remote, [2, 6], b"compact") + + registration = endpoint.registrations[0] + assert registration.region_calls[-1][4] == {"stride": 64, "count": 8} + assert endpoint.peers[0].region_calls[-1][3] == { + "device": "cuda:0", + "memory_type": "VRAM", + "stride": 64, + "count": 8, + } + assert endpoint.plans[-1].kwargs["local"] == local.regions + assert endpoint.plans[-1].kwargs["remote"] == remote.regions + assert endpoint.submissions == [] + assert agent.transfer(handle) == "PROC" + assert tuple(endpoint.submissions[-1][2]["local_indices"]) == (7, 1) + assert tuple(endpoint.submissions[-1][2]["remote_indices"]) == (2, 6) + assert handle.work.test_calls == 0 + handle.work.done = True + handle.work.state = _WorkState.SUCCEEDED + agent.progress() + assert handle.work.test_calls == 0 + assert agent.check_xfer_state(handle) == "DONE" + assert handle.work.test_calls == 1 + assert handle.work.closed is True + + +def test_status_observation_failure_stays_nonterminal(): + agent, _, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") + assert agent.transfer(handle) == "PROC" + handle.work.test_error = RuntimeError("temporary status failure") + + assert agent.check_xfer_state(handle) == "PROC" + assert handle.terminal_state is None + assert handle.work.closed is False + + handle.work.test_error = None + handle.work.done = True + handle.work.state = _WorkState.SUCCEEDED + agent.progress() + assert handle.work.test_calls == 2 + assert agent.check_xfer_state(handle) == "DONE" + assert handle.terminal_state == "DONE" + assert handle.work.closed is True + + +def test_batch_progress_leaves_one_status_probe_per_active_work(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + handles = [ + agent.initialize_xfer("WRITE", local, remote, "decode-agent") for _ in range(3) + ] + assert [agent.transfer(handle) for handle in handles] == ["PROC"] * 3 + handles[0].work.done = True + handles[0].work.state = _WorkState.SUCCEEDED + + agent.progress() + assert [handle.work.test_calls for handle in handles] == [1, 1, 1] + states = [agent.check_xfer_state(handle) for handle in handles] + + assert endpoint.progress_calls == 1 + assert states == ["DONE", "PROC", "PROC"] + assert [handle.work.test_calls for handle in handles] == [2, 2, 2] + assert handles[0].work.closed is True + assert handles[0].owned_plan.closed is True + assert handles[1].work.closed is False + assert handles[1].owned_plan.closed is False + assert handles[2].work.closed is False + assert handles[2].owned_plan.closed is False + assert set(agent._works) == {id(handles[1]), id(handles[2])} + + handles[1].work.done = True + handles[1].work.state = _WorkState.SUCCEEDED + agent.progress() + assert [handle.work.test_calls for handle in handles] == [2, 2, 2] + states = [agent.check_xfer_state(handle) for handle in handles[1:]] + + assert endpoint.progress_calls == 2 + assert states == ["DONE", "PROC"] + assert [handle.work.test_calls for handle in handles] == [2, 3, 3] + assert handles[1].work.closed is True + assert handles[1].owned_plan.closed is True + assert handles[2].work.closed is False + assert handles[2].owned_plan.closed is False + assert set(agent._works) == {id(handles[2])} + + +def test_detached_notification_work_is_still_reaped_by_progress(): + agent, endpoint, _ = _make_agent(synchronous_notification_send=False) + _register_and_add_peer(agent) + handle = agent.send_notif("decode-agent", b"abort") + assert endpoint.notifications[-1][2] is handle + assert endpoint.open_works_reads == 0 + handle.work.done = True + handle.work.state = _WorkState.SUCCEEDED + + agent.progress() + + assert endpoint.progress_calls == 1 + assert handle.work.test_calls == 1 + assert handle.work.closed is True + assert handle.terminal_state == "DONE" + assert agent._works == {} + assert agent._detached_works == [] + + +def test_uncancellable_sibling_is_retained_then_reaped_when_terminal(): + agent, _, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + failed, sibling, foreground = [ + agent.initialize_xfer("WRITE", local, remote, "decode-agent") for _ in range(3) + ] + assert [agent.transfer(handle) for handle in (failed, sibling, foreground)] == [ + "PROC", + "PROC", + "PROC", + ] + failed.work.done = True + failed.work.state = _WorkState.FAILED + sibling.work.cancel_completes = False + + agent.cancel_handles([failed, sibling]) + + assert failed.work.closed is True + assert sibling.work.closed is False + assert id(sibling) in agent._works + assert id(sibling) in agent._retained_work_ids + sweep_thread = agent._sweep_thread + assert sweep_thread is not None + foreground_probe_count = foreground.work.test_calls + + sibling.work.done = True + sibling.work.state = _WorkState.SUCCEEDED + deadline = time.monotonic() + 1 + while not sibling.work.closed and time.monotonic() < deadline: + time.sleep(0.01) + + assert sibling.work.closed is True + sweep_thread.join(timeout=1) + assert sweep_thread.is_alive() is False + assert agent._retained_work_ids == set() + assert foreground.work.test_calls == foreground_probe_count + foreground.work.done = True + foreground.work.state = _WorkState.SUCCEEDED + assert agent.check_xfer_state(foreground) == "DONE" + + +def test_cancel_handles_and_wait_does_not_return_before_uncancellable_work_is_safe(): + agent, _, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + failed, sibling = [ + agent.initialize_xfer("WRITE", local, remote, "decode-agent") for _ in range(2) + ] + assert agent.transfer(failed) == "PROC" + assert agent.transfer(sibling) == "PROC" + failed.work.done = True + failed.work.state = _WorkState.FAILED + sibling.work.cancel_completes = False + + def complete_sibling() -> None: + time.sleep(0.01) + sibling.work.done = True + sibling.work.state = _WorkState.SUCCEEDED + + completion = threading.Thread(target=complete_sibling) + completion.start() + agent.cancel_handles_and_wait([failed, sibling]) + completion.join(timeout=1) + + assert completion.is_alive() is False + assert failed.work.closed is True + assert sibling.work.closed is True + assert agent._works == {} + + +def test_background_progress_does_not_call_manual_progress(): + agent, endpoint, _ = _make_agent(background_progress=True) + _register_and_add_peer(agent) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") + + class _UnexpectedLock: + def __enter__(self): + raise AssertionError("empty background progress acquired the lock") + + def __exit__(self, *args): + return False + + with patch.object(agent, "_lock", _UnexpectedLock()): + agent.progress() + + assert endpoint.progress_calls == 0 + with pytest.raises(RuntimeError, match="has not been posted"): + agent.check_xfer_state(handle) + assert endpoint.submissions == [] + assert agent.transfer(handle) == "PROC" + assert handle.work.test_calls == 1 + + handle.work.done = True + handle.work.state = _WorkState.SUCCEEDED + agent.progress() + + assert endpoint.progress_calls == 0 + assert handle.work.test_calls == 1 + assert agent.check_xfer_state(handle) == "DONE" + assert handle.work.test_calls == 2 + assert handle.work.closed is True + assert handle.owned_plan.closed is True + + +def test_multiple_mode_admission_gate_drains_progress_before_close(): + agent, endpoint, _ = _make_agent(background_progress=False) + progress_entered = threading.Event() + progress_release = threading.Event() + close_started = threading.Event() + close_finished = threading.Event() + failures: list[BaseException] = [] + original_progress = endpoint.progress + + def blocking_progress(): + progress_entered.set() + if not progress_release.wait(1): + raise RuntimeError("test progress release timed out") + original_progress() + + endpoint.progress = blocking_progress + + def run_progress(): + try: + agent.progress() + except BaseException as error: + failures.append(error) + + def run_close(): + close_started.set() + try: + agent.close() + except BaseException as error: + failures.append(error) + finally: + close_finished.set() + + progress_thread = threading.Thread(target=run_progress) + close_thread = threading.Thread(target=run_close) + progress_thread.start() + assert progress_entered.wait(1) + close_thread.start() + assert close_started.wait(1) + assert not close_finished.wait(0.05) + progress_release.set() + progress_thread.join(timeout=1) + close_thread.join(timeout=1) + + assert not progress_thread.is_alive() + assert not close_thread.is_alive() + assert failures == [] + assert endpoint.closed is True + + +def test_distinct_handle_submits_overlap_without_adapter_locks(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") + remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") + handles = [ + agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) for _ in range(2) + ] + barrier = threading.Barrier(2) + counter_lock = threading.Lock() + active = 0 + max_active = 0 + original_submit = endpoint._submit + + def overlapping_submit(operation, plan, kwargs): + nonlocal active, max_active + assert not agent._lock._is_owned() + assert not kwargs["recovery_token"]._state_lock.locked() + with counter_lock: + active += 1 + max_active = max(max_active, active) + try: + barrier.wait(timeout=1) + return original_submit(operation, plan, kwargs) + finally: + with counter_lock: + active -= 1 + + endpoint._submit = overlapping_submit + results = [] + failures = [] + + def submit(handle): + try: + results.append(agent.transfer(handle)) + except BaseException as error: + failures.append(error) + + threads = [threading.Thread(target=submit, args=(handle,)) for handle in handles] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=2) + + assert all(not thread.is_alive() for thread in threads) + assert failures == [] + assert sorted(results) == ["PROC", "PROC"] + assert max_active == 2 + + +def test_distinct_handle_status_refreshes_overlap_without_adapter_locks(): + agent, _, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + handles = [ + agent.initialize_xfer("WRITE", local, remote, "decode-agent") for _ in range(2) + ] + assert [agent.transfer(handle) for handle in handles] == ["PROC", "PROC"] + barrier = threading.Barrier(2) + counter_lock = threading.Lock() + active = 0 + max_active = 0 + + def make_hook(handle): + def hook(): + nonlocal active, max_active + assert not agent._lock._is_owned() + assert not handle._state_lock.locked() + with counter_lock: + active += 1 + max_active = max(max_active, active) + try: + barrier.wait(timeout=1) + finally: + with counter_lock: + active -= 1 + + return hook + + for handle in handles: + handle.work.state_hook = make_hook(handle) + results = [] + failures = [] + + def poll(handle): + try: + results.append(agent.check_xfer_state(handle)) + except BaseException as error: + failures.append(error) + + threads = [threading.Thread(target=poll, args=(handle,)) for handle in handles] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=2) + + assert all(not thread.is_alive() for thread in threads) + assert failures == [] + assert sorted(results) == ["PROC", "PROC"] + assert max_active == 2 + + +def test_teardown_gate_never_refreshes_status_under_agent_lock(): + agent, _, _ = _make_agent() + registrations = _register_and_add_peer(agent) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") + assert agent.transfer(handle) == "PROC" + with handle._state_lock: + handle.submission_interrupted = True + probes_before = handle.work.test_calls + open_works_before = agent._endpoint.open_works_reads + handle.work.state_hook = lambda: pytest.fail( + "teardown gate must not perform a Core status read" + ) + + with pytest.raises(RuntimeError, match="submission.*unsettled"): + agent.deregister_memory(registrations) + + assert handle.work.test_calls == probes_before + assert agent._endpoint.open_works_reads == open_works_before + + +def test_concurrent_interrupted_recovery_has_one_core_snapshot_owner(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + handle = _WorkHandle(phase="recovery_pending") + handle.submission_interrupted = True + agent._track_work(handle) + recovered = _Work(handle) + endpoint._open_works.append(recovered) + entered = threading.Event() + release = threading.Event() + original_open_works = type(endpoint).open_works.fget + + def blocking_open_works(instance): + entered.set() + assert release.wait(1) + return original_open_works(instance) + + results = [] + failures = [] + + def recover(): + try: + results.append(agent._recover_interrupted_post_work(handle)) + except BaseException as error: + failures.append(error) + + with patch.object(type(endpoint), "open_works", property(blocking_open_works)): + owner = threading.Thread(target=recover) + contender = threading.Thread(target=recover) + owner.start() + assert entered.wait(1) + contender.start() + contender.join(timeout=1) + release.set() + owner.join(timeout=1) + + assert not owner.is_alive() + assert not contender.is_alive() + assert failures == [] + assert sorted(results) == ["pending", "recovered"] + assert endpoint.open_works_reads == 1 + assert handle.work is recovered + assert handle.phase == "posted" + + +def test_manual_progress_does_not_hold_adapter_state_lock(): + agent, endpoint, _ = _make_agent(background_progress=False) + observed = [] + + def inspect_progress_lock(): + observed.append(agent._lock._is_owned()) + + endpoint.progress = inspect_progress_lock + agent.progress() + assert observed == [False] + + +def test_close_adopts_authoritative_closed_state_before_child_queries(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + snapshots_before = endpoint.construction_snapshot_reads + endpoint.poison_construction_snapshots = True + endpoint.closed = True + + agent.close() + + assert endpoint.construction_snapshot_reads == snapshots_before + assert agent._endpoint is None + assert agent._works == {} + assert agent._peers == {} + + +def test_close_commit_lost_return_converges_on_retry(): + agent, endpoint, _ = _make_agent() + + def committed_then_interrupted(): + endpoint.close_calls += 1 + endpoint.closed = True + raise KeyboardInterrupt("close return lost") + + endpoint.close = committed_then_interrupted + with pytest.raises(KeyboardInterrupt, match="close return lost"): + agent.close() + + assert agent._endpoint is None + assert endpoint.close_calls == 1 + agent.close() + assert endpoint.close_calls == 1 + + +def test_background_progress_reaps_nonempty_detached_work(): + agent, endpoint, _ = _make_agent( + background_progress=True, synchronous_notification_send=False + ) + _register_and_add_peer(agent) + handle = agent.send_notif("decode-agent", b"abort") + handle.work.done = True + handle.work.state = _WorkState.SUCCEEDED + + agent.progress() + + assert endpoint.progress_calls == 0 + assert handle.work.test_calls == 1 + assert handle.work.closed is True + assert agent._detached_works == [] + + +def test_conn_enables_background_progress_for_prefill_and_decode_roles(): + repo_root = Path(__file__).resolve().parents[4] + source = (repo_root / "python/sglang/srt/disaggregation/nixl/conn.py").read_text() + tree = ast.parse(source) + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "TorchTransferAgent" + ] + assert len(calls) == 1 + keywords = {keyword.arg: keyword.value for keyword in calls[0].keywords} + assert isinstance(keywords["background_progress"], ast.Constant) + assert keywords["background_progress"].value is True + assert isinstance(keywords["num_threads"], ast.Name) + assert keywords["num_threads"].id == "num_threads" + + assignments = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "num_threads" + for target in node.targets + ) + ] + role_assignment = next( + node for node in assignments if isinstance(node.value, ast.IfExp) + ) + assert isinstance(role_assignment.value.body, ast.Constant) + assert role_assignment.value.body.value == 8 + assert isinstance(role_assignment.value.orelse, ast.Constant) + assert role_assignment.value.orelse.value == 0 + + +def test_registration_names_do_not_collide_after_deregister(): + agent, endpoint, _ = _make_agent() + first = agent.register_memory([(0x1000, 0x100, 0, "")], "VRAM") + second = agent.register_memory([(0x2000, 0x100, 0, "")], "VRAM") + agent.deregister_memory(first) + third = agent.register_memory([(0x3000, 0x100, 0, "")], "VRAM") + + assert second[0].name == "sglang_region_1" + assert third[0].name == "sglang_region_2" + assert len({registration.name for registration in endpoint.registrations}) == 3 + assert agent._registration_groups == [second, third] + + +@pytest.mark.parametrize("operation", ["READ", "WRITE"]) +def test_one_off_transfer_orients_local_and_remote_regions(operation): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.get_xfer_descs([(0x1080, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5080, 32, 0)], "VRAM") + source, destination = (remote, local) if operation == "READ" else (local, remote) + + handle = agent.initialize_xfer( + operation, source, destination, "decode-agent", b"done" + ) + + plan = endpoint.plans[-1] + assert plan.kwargs["indexed"] is False + assert plan.kwargs["local"][0][0] == "local" + assert plan.kwargs["remote"][0][0] == "remote" + assert endpoint.submissions == [] + assert agent.transfer(handle) == "PROC" + assert endpoint.submissions[-1][0] == operation + assert endpoint.submissions[-1][2]["notification"] == b"done" + assert endpoint.submissions[-1][2]["recovery_token"] is handle + assert endpoint.open_works_reads == 0 + handle.work.done = True + handle.work.state = _WorkState.SUCCEEDED + assert agent.check_xfer_state(handle) == "DONE" + assert plan.closed is True + + +def test_standalone_notifications_preserve_peer_alias(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + notification = SimpleNamespace( + source_endpoint_id="remote-id", + source_incarnation="remote-incarnation", + payload=b"9_aux", + ) + endpoint.incoming_notifications.append(notification) + + payload = b"abort" + result = agent.send_notif("decode-agent", payload) + + assert agent.get_new_notifs() == {"decode-agent": [b"9_aux"]} + assert endpoint.progress_calls == 0 + assert result is None + assert endpoint.synchronous_notifications[0][1] is payload + assert endpoint.notifications == [] + assert agent._works == {} + assert agent._detached_works == [] + + +def test_notification_without_imported_source_uses_explicit_sentinel(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + endpoint.incoming_notifications.append( + SimpleNamespace( + source_endpoint_id=None, + source_incarnation=None, + payload=b"9_aux", + ) + ) + + assert agent.get_new_notifs() == {"": [b"9_aux"]} + + +def test_notification_from_stale_incarnation_is_rejected(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + endpoint.incoming_notifications.append( + SimpleNamespace( + source_endpoint_id="remote-id", + source_incarnation="old-incarnation", + payload=b"9_aux", + ) + ) + + with pytest.raises(RuntimeError, match="unknown or stale"): + agent.get_new_notifs() + + +def test_notification_batches_group_payloads_and_accept_anonymous_source(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + + class CountingAliases(dict): + lookups = 0 + + def __getitem__(self, key): + self.lookups += 1 + return super().__getitem__(key) + + aliases = CountingAliases(agent._source_aliases) + agent._source_aliases = aliases + endpoint.poll_notification_batches = lambda: [ + SimpleNamespace( + source_endpoint_id="remote-id", + source_incarnation="remote-incarnation", + payloads=(b"first", b"second"), + ), + SimpleNamespace( + source_endpoint_id=None, + source_incarnation=None, + payloads=(b"anonymous",), + ), + ] + + def reject_scalar_poll(): + raise AssertionError("grouped receive must not retry scalar polling") + + endpoint.poll_notifications = reject_scalar_poll + + assert agent.get_new_notifs() == { + "decode-agent": [b"first", b"second"], + "": [b"anonymous"], + } + assert aliases.lookups == 1 + + +def test_notification_batch_from_stale_incarnation_is_rejected(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + endpoint.poll_notification_batches = lambda: [ + SimpleNamespace( + source_endpoint_id="remote-id", + source_incarnation="old-incarnation", + payloads=(b"stale",), + ) + ] + + with pytest.raises(RuntimeError, match="unknown or stale"): + agent.get_new_notifs() + + +def test_notification_batch_failure_is_never_retried_as_scalar(): + agent, endpoint, _ = _make_agent() + batch_calls = 0 + scalar_calls = 0 + + def fail_batch_poll(): + nonlocal batch_calls + batch_calls += 1 + raise RuntimeError("destructive batch failure") + + def count_scalar_poll(): + nonlocal scalar_calls + scalar_calls += 1 + return [] + + endpoint.poll_notification_batches = fail_batch_poll + endpoint.poll_notifications = count_scalar_poll + + with pytest.raises(RuntimeError, match="destructive batch failure"): + agent.get_new_notifs() + assert batch_calls == 1 + assert scalar_calls == 0 + + +def test_active_work_prevents_unsafe_close(): + agent, _, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") + assert agent.transfer(handle) == "PROC" + handle.work.cancel_completes = False + + with pytest.raises(RuntimeError, match="active transfer"): + agent.close() + + handle.work.done = True + handle.work.state = _WorkState.CANCELLED + agent.close() + + +def test_prepared_release_and_agent_close_never_post(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + local_prepared = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") + remote_prepared = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") + indexed = agent.make_prepped_xfer( + "WRITE", local_prepared, [0], remote_prepared, [0], b"never-posted" + ) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + owned = agent.initialize_xfer( + "WRITE", local, remote, "decode-agent", b"also-never-posted" + ) + + agent.release_xfer_handle(owned) + assert owned.owned_plan.closed is True + assert endpoint.submit_attempts == 0 + assert indexed.local_indices == (0,) + assert indexed.remote_indices == (0,) + agent.close() + + assert indexed.phase == "released" + assert indexed.local_indices is None + assert indexed.remote_indices is None + assert endpoint.submit_attempts == 0 + assert all(plan.closed for plan in endpoint.plans) + + +def test_post_failure_is_terminal_and_never_retried(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") + remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") + handle = agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) + endpoint.submit_error = RuntimeError("post failed") + + with pytest.raises(RuntimeError, match="post failed"): + agent.transfer(handle) + assert endpoint.submit_attempts == 1 + assert handle.terminal_state == "ERR" + assert handle.phase == "released" + assert agent.transfer(handle) == "ERR" + assert endpoint.submit_attempts == 1 + + +@pytest.mark.parametrize("indexed", [True, False], ids=["indexed", "one-off"]) +def test_postcommit_base_exception_recovers_exact_core_work(indexed): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + endpoint.postcommit_error = KeyboardInterrupt("submit return interrupted") + + if indexed: + endpoint.next_fused_state = _WorkState.COMPLETED + local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") + remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") + handle = agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) + else: + endpoint.next_work_done = True + endpoint.next_work_state = _WorkState.COMPLETED + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") + + with pytest.raises(KeyboardInterrupt, match="submit return interrupted"): + agent.transfer(handle) + + recovered = endpoint.submissions[-1][3] + assert handle.work is recovered + assert recovered.close_calls == 1 + assert recovered.closed is True + assert endpoint.open_works == () + assert endpoint.open_works_reads == 2 + assert endpoint.submit_attempts == 1 + assert endpoint.submissions[-1][2]["recovery_token"] is handle + assert handle.terminal_state == "ERR" + assert handle.phase == "released" + assert id(handle) not in agent._works + assert agent.transfer(handle) == "ERR" + assert endpoint.submit_attempts == 1 + if indexed: + assert endpoint.plans[-1].closed is False + else: + assert handle.owned_plan.close_calls == 1 + assert handle.owned_plan.closed is True + + +class _IdentityBombHandle(_WorkHandle): + def __hash__(self): + raise AssertionError("recovery token was hashed") + + def __eq__(self, other): + raise AssertionError("recovery token equality was evaluated") + + def __bool__(self): + raise AssertionError("recovery token truthiness was evaluated") + + +class _UnprintableInterrupt(KeyboardInterrupt): + def __str__(self): + raise RuntimeError("exception formatting is hostile") + + +def test_interrupted_post_ignores_distractor_and_uses_identity_only(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + distractor = _Work(object()) + endpoint._open_works.append(distractor) + endpoint.next_work_done = True + endpoint.next_work_state = _WorkState.COMPLETED + endpoint.postcommit_error = KeyboardInterrupt("submit return interrupted") + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + with patch( + "sglang.srt.disaggregation.nixl.torch_transfer._WorkHandle", + _IdentityBombHandle, + ): + handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") + + with pytest.raises(KeyboardInterrupt, match="submit return interrupted"): + agent.transfer(handle) + + recovered = endpoint.submissions[-1][3] + assert endpoint.submissions[-1][2]["recovery_token"] is handle + assert recovered.closed is True + assert distractor.closed is False + assert handle.terminal_state == "ERR" + assert handle.phase == "released" + + +def test_ambiguous_token_recovery_is_retained_and_retryable(): + agent, endpoint, _ = _make_agent() + registrations = _register_and_add_peer(agent) + endpoint.postcommit_error = KeyboardInterrupt("submit return interrupted") + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") + duplicate = _Work(handle) + + def add_duplicate(): + endpoint._open_works.append(duplicate) + + endpoint.postcommit_hook = add_duplicate + with patch.object(agent, "_start_retained_work_sweeper") as start_sweeper: + assert agent.transfer(handle) == "PROC" + + assert handle.phase == "recovery_pending" + assert handle.work is None + assert id(handle) in agent._works + assert id(handle) in agent._retained_work_ids + assert start_sweeper.call_count == 1 + assert agent.transfer(handle) == "PROC" + with pytest.raises(RuntimeError, match="recovery is pending"): + agent.release_xfer_handle(handle) + with pytest.raises(RuntimeError, match="submission.*unsettled"): + agent.deregister_memory(registrations) + with pytest.raises(RuntimeError, match="submission.*unsettled"): + agent.remove_remote_agent("decode-agent") + with pytest.raises(RuntimeError, match="active transfer"): + agent.close() + + duplicate.close() + assert agent.transfer(handle) == "ERR" + assert endpoint.submit_attempts == 1 + assert handle.phase == "released" + assert endpoint.submissions[-1][3].closed is True + agent.remove_remote_agent("decode-agent") + agent.deregister_memory(registrations) + + +def test_interrupted_open_works_lookup_retries_without_resubmission(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + endpoint.postcommit_error = KeyboardInterrupt("submit return interrupted") + endpoint.open_works_errors.append(RuntimeError("snapshot interrupted")) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") + + with patch.object(agent, "_start_retained_work_sweeper"): + assert agent.transfer(handle) == "PROC" + assert handle.phase == "recovery_pending" + assert agent.transfer(handle) == "ERR" + + assert endpoint.submit_attempts == 1 + assert endpoint.open_works_reads == 2 + assert handle.phase == "released" + assert "KeyboardInterrupt: submit return interrupted" in str( + agent.pop_xfer_error(handle) + ) + + +def test_exact_recovery_keeps_teardown_blocked_until_work_is_terminal(): + agent, endpoint, _ = _make_agent() + registrations = _register_and_add_peer(agent) + endpoint.postcommit_error = KeyboardInterrupt("submit return interrupted") + + def make_recovered_work_uncancellable(): + endpoint.submissions[-1][3].cancel_completes = False + + endpoint.postcommit_hook = make_recovered_work_uncancellable + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") + + with patch.object(agent, "_start_retained_work_sweeper"): + assert agent.transfer(handle) == "PROC" + assert handle.phase == "posted" + assert handle.work.closed is False + with pytest.raises(RuntimeError, match="submission.*unsettled"): + agent.deregister_memory(registrations) + with pytest.raises(RuntimeError, match="submission.*unsettled"): + agent.remove_remote_agent("decode-agent") + + handle.work.done = True + handle.work.state = _WorkState.COMPLETED + assert agent.transfer(handle) == "ERR" + + assert endpoint.submit_attempts == 1 + assert handle.work.closed is True + assert "KeyboardInterrupt: submit return interrupted" in str( + agent.pop_xfer_error(handle) + ) + agent.remove_remote_agent("decode-agent") + agent.deregister_memory(registrations) + + +def test_recovery_precedes_hostile_exception_diagnostics_and_batch_owns_handle(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + endpoint.postcommit_error = _UnprintableInterrupt() + + def make_recovered_work_uncancellable(): + endpoint.submissions[-1][3].cancel_completes = False + + endpoint.postcommit_hook = make_recovered_work_uncancellable + handles = [] + agent.begin_handle_batch(handles) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") + + with patch.object(agent, "_start_retained_work_sweeper"): + assert handles == [handle] + assert agent.transfer(handle) == "PROC" + assert endpoint.submit_attempts == 1 + assert handle.work.closed is False + + handle.work.done = True + handle.work.state = _WorkState.COMPLETED + assert agent.check_xfer_state(handle) == "ERR" + + agent.end_handle_batch(handles) + assert handle.work.closed is True + assert "diagnostic formatting unavailable" in str(agent.pop_xfer_error(handle)) + + +def test_async_notification_lost_return_recovers_by_token(): + agent, endpoint, _ = _make_agent(synchronous_notification_send=False) + _register_and_add_peer(agent) + distractor = _Work(object()) + endpoint._open_works.append(distractor) + endpoint.postcommit_error = KeyboardInterrupt("notification return interrupted") + + with pytest.raises(KeyboardInterrupt, match="notification return interrupted"): + agent.send_notif("decode-agent", b"abort") + + _, payload, token, work = endpoint.notifications[-1] + assert payload == b"abort" + assert token.work is work + assert token.terminal_state == "ERR" + assert token.phase == "released" + assert work.closed is True + assert distractor.closed is False + assert endpoint.open_works_reads == 1 + assert agent._works == {} + assert agent._detached_works == [] + + +def test_async_notification_tracking_interruption_rolls_back_hidden_handle(): + agent, endpoint, _ = _make_agent(synchronous_notification_send=False) + _register_and_add_peer(agent) + original_track = agent._track_work + + def store_then_interrupt(handle): + original_track(handle) + raise KeyboardInterrupt("tracking return interrupted") + + with patch.object(agent, "_track_work", side_effect=store_then_interrupt): + with pytest.raises(KeyboardInterrupt, match="tracking return interrupted"): + agent.send_notif("decode-agent", b"abort") + + assert endpoint.notifications == [] + assert agent._works == {} + assert agent._detached_works == [] + + +def test_prepared_plan_teardown_quarantines_then_rejects_stale_post(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") + remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") + handle = agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) + plan = endpoint.plans[-1] + + with pytest.raises(RuntimeError, match="referenced by a transfer handle"): + agent.remove_remote_agent("decode-agent") + assert plan.closed is False + assert "decode-agent" in agent._retiring_peers + assert endpoint.submit_attempts == 0 + + with pytest.raises(RuntimeError, match="teardown is in progress"): + agent.transfer(handle) + assert endpoint.submit_attempts == 0 + assert endpoint.open_works_reads == 0 + assert handle.terminal_state == "ERR" + agent.remove_remote_agent("decode-agent") + assert plan.closed is True + assert "decode-agent" not in agent._retiring_peers + + +@pytest.mark.parametrize("version", [None, True, 1]) +def test_core_factory_v2_gate_rejects_absent_bool_and_old_versions(version): + api = _make_api() + api.BACKEND_FACTORY_API_VERSION = version + + with patch( + "sglang.srt.disaggregation.nixl.torch_transfer.importlib.import_module", + return_value=api, + ): + with pytest.raises(ImportError, match="FactoryV2"): + _load_transfer_api() + + +def test_core_factory_v2_gate_requires_marker_and_accepts_newer_core(): + api = _make_api() + del api.BackendFactoryV2 + with patch( + "sglang.srt.disaggregation.nixl.torch_transfer.importlib.import_module", + return_value=api, + ): + with pytest.raises(ImportError, match="FactoryV2"): + _load_transfer_api() + + api = _make_api() + api.BACKEND_FACTORY_API_VERSION = 3 + with patch( + "sglang.srt.disaggregation.nixl.torch_transfer.importlib.import_module", + return_value=api, + ): + assert _load_transfer_api() is api + + +@pytest.mark.parametrize("version", [None, True, 1, 3]) +def test_provider_factory_v2_gate_requires_exact_integer_abi(version): + register = MagicMock() + provider = SimpleNamespace( + TORCH_TRANSFER_FACTORY_API_VERSION=version, + register_torch_backend=register, + ) + with patch( + "sglang.srt.disaggregation.nixl.torch_transfer.importlib.import_module", + return_value=provider, + ): + with pytest.raises(ImportError, match="exact FactoryV2"): + _register_nixl_backend() + register.assert_not_called() + + provider.TORCH_TRANSFER_FACTORY_API_VERSION = 2 + with patch( + "sglang.srt.disaggregation.nixl.torch_transfer.importlib.import_module", + return_value=provider, + ): + _register_nixl_backend() + register.assert_called_once_with() + + +@pytest.mark.parametrize("mismatch", [False, True], ids=["exact", "foreign"]) +def test_constructor_recovery_is_identity_gated_and_cleanup_is_retryable(mismatch): + class HostileCleanup(RuntimeError): + def __str__(self): + raise AssertionError("cleanup diagnostics evaluated __str__") + + agent = _make_uncreated_agent() + error = KeyboardInterrupt("constructor cut") + _Endpoint.constructor_cut = ( + error, + mismatch, + ( + [AssertionError("mismatched endpoint must not close")] + if mismatch + else [HostileCleanup()] + ), + ) + + with pytest.raises(KeyboardInterrupt, match="constructor cut"): + agent.create_backend("UCX", {}) + + candidate = _Endpoint.instances[-1] + assert agent._endpoint is None + if mismatch: + assert candidate.close_calls == 0 + assert candidate.closed is False + assert agent._endpoint_recovery is None + assert any("non-matching" in note for note in error.__notes__) + candidate.close_errors.clear() + candidate.close() + else: + assert candidate.close_calls == 1 + assert candidate.closed is False + assert agent._endpoint_recovery is candidate + assert any( + "cleanup also raised:" in note and "HostileCleanup" in note + for note in error.__notes__ + ) + agent.create_backend("UCX", {}) + assert candidate.close_calls == 2 + assert candidate.closed is True + assert agent._endpoint_recovery is None + + +def test_postpublication_constructor_cut_rolls_back_all_publications(): + agent = _make_uncreated_agent() + armed = True + + def store_then_interrupt(self, name, value): + nonlocal armed + object.__setattr__(self, name, value) + if self is agent and name == "_endpoint" and type(value) is _Endpoint and armed: + armed = False + value.close_errors.append(RuntimeError("close cut")) + raise KeyboardInterrupt("after endpoint publication") + + with patch.object(TorchTransferAgent, "__setattr__", store_then_interrupt): + with pytest.raises(KeyboardInterrupt, match="after endpoint publication"): + agent.create_backend("UCX", {}) + + failed_endpoint = _Endpoint.instances[-1] + assert agent._endpoint is None + assert agent._transport_backend is None + assert agent._send_notification_sync is None + assert agent._endpoint_recovery is failed_endpoint + agent.create_backend("UCX", {}) + assert failed_endpoint.closed is True + assert agent._endpoint is _Endpoint.instances[-1] + assert agent._endpoint is not failed_endpoint + + +def test_capability_failure_rolls_back_endpoint_before_publication(): + agent = _make_uncreated_agent() + _Endpoint.missing_capability = "write" + + with pytest.raises(NotImplementedError, match="write"): + agent.create_backend("UCX", {}) + + failed_endpoint = _Endpoint.instances[-1] + assert failed_endpoint.closed is True + assert agent._endpoint is None + assert agent._transport_backend is None + assert agent._send_notification_sync is None + + +def test_registration_and_peer_lost_returns_close_only_exact_orphans(): + agent, endpoint, _ = _make_agent() + known_registration = agent.register_memory([(0x1000, 0x100, 0, "")], "VRAM") + original_register = endpoint.register + orphan_registration = None + + def register_then_interrupt(*args, **kwargs): + nonlocal orphan_registration + orphan_registration = original_register(*args, **kwargs) + raise KeyboardInterrupt("registration return cut") + + with patch.object(endpoint, "register", side_effect=register_then_interrupt): + with pytest.raises(KeyboardInterrupt, match="registration return cut"): + agent.register_memory([(0x2000, 0x100, 0, "")], "VRAM") + + assert orphan_registration.closed is True + assert known_registration[0].closed is False + assert endpoint.open_registrations == known_registration + assert agent._construction_dirty is False + + original_import = endpoint.import_peer + orphan_peer = None + + def import_then_interrupt(*args, **kwargs): + nonlocal orphan_peer + orphan_peer = original_import(*args, **kwargs) + raise KeyboardInterrupt("peer return cut") + + with patch.object(endpoint, "import_peer", side_effect=import_then_interrupt): + with pytest.raises(KeyboardInterrupt, match="peer return cut"): + agent.add_remote_agent(b"peer-metadata", peer_name="decode-agent") + + assert orphan_peer.closed is True + assert endpoint.open_peers == () + assert agent._peers == {} + assert agent._construction_dirty is False + + +@pytest.mark.parametrize("indexed", [True, False], ids=["cached", "one-off"]) +def test_plan_lost_return_preserves_known_child_and_closes_exact_orphan(indexed): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + known_local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") + known_remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") + agent.make_prepped_xfer("WRITE", known_local, [0], known_remote, [0]) + known_plan = endpoint.plans[-1] + original_prepare = endpoint.prepare + orphan_plan = None + + def prepare_then_interrupt(**kwargs): + nonlocal orphan_plan + orphan_plan = original_prepare(**kwargs) + raise KeyboardInterrupt("plan return cut") + + with patch.object(endpoint, "prepare", side_effect=prepare_then_interrupt): + with pytest.raises(KeyboardInterrupt, match="plan return cut"): + if indexed: + local = agent.prep_xfer_dlist("", [(0x1040, 32, 0)], "VRAM") + remote = agent.prep_xfer_dlist( + "decode-agent", [(0x5040, 32, 0)], "VRAM" + ) + agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) + else: + local = agent.get_xfer_descs([(0x1040, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5040, 32, 0)], "VRAM") + agent.initialize_xfer("WRITE", local, remote, "decode-agent") + + assert orphan_plan.closed is True + assert known_plan.closed is False + assert endpoint.open_plans == (known_plan,) + assert agent._construction_dirty is False + + +def test_dirty_plan_cleanup_cut_is_retried_before_next_construction(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + original_prepare = endpoint.prepare + orphan_plan = None + + def prepare_then_interrupt(**kwargs): + nonlocal orphan_plan + orphan_plan = original_prepare(**kwargs) + orphan_plan.close_errors.append(RuntimeError("first close cut")) + raise KeyboardInterrupt("plan return cut") + + with patch.object(endpoint, "prepare", side_effect=prepare_then_interrupt): + with pytest.raises(KeyboardInterrupt) as caught: + agent.initialize_xfer("WRITE", local, remote, "decode-agent") + + assert orphan_plan.closed is False + assert orphan_plan.close_calls == 1 + assert agent._construction_dirty is True + assert any("cleanup also raised" in note for note in caught.value.__notes__) + + handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") + assert orphan_plan.closed is True + assert orphan_plan.close_calls == 2 + assert handle.owned_plan is endpoint.plans[-1] + assert agent._construction_dirty is False + + +def test_reconcile_scans_future_slot_and_binding_children_first(): + agent, endpoint, _ = _make_agent() + close_order = [] + + class OrderedResource(_Resource): + def __init__(self, name): + super().__init__() + self.name = name + + def close(self): + close_order.append(self.name) + super().close() + + slot = OrderedResource("slot") + binding = OrderedResource("binding") + plan = OrderedResource("plan") + peer = OrderedResource("peer") + registration = OrderedResource("registration") + endpoint.transfer_slots.append(slot) + endpoint.bound_transfers.append(binding) + endpoint.plans.append(plan) + endpoint.peers.append(peer) + endpoint.registrations.append(registration) + agent._construction_dirty = True + + agent.register_memory([(0x1000, 0x100, 0, "")], "VRAM") + + assert slot.closed is True + assert binding.closed is True + assert plan.closed is True + assert peer.closed is True + assert registration.closed is True + assert close_order == ["slot", "binding", "plan", "peer", "registration"] + assert agent._construction_dirty is False + + +def test_clean_one_off_and_cached_hit_never_read_construction_snapshots(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + endpoint.construction_snapshot_reads = 0 + endpoint.poison_construction_snapshots = True + local_raw = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote_raw = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + agent.initialize_xfer("WRITE", local_raw, remote_raw, "decode-agent") + local = agent.prep_xfer_dlist("", [(0x1040, 32, 0)], "VRAM") + remote = agent.prep_xfer_dlist("decode-agent", [(0x5040, 32, 0)], "VRAM") + agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) + agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) + + assert endpoint.construction_snapshot_reads == 0 + + +def test_peer_postpublication_cut_retries_without_duplicate_core_child(): + class StoreThenInterrupt(dict): + def __setitem__(self, key, value): + super().__setitem__(key, value) + raise KeyboardInterrupt("peer map STORE cut") + + agent, endpoint, _ = _make_agent() + agent._peers = StoreThenInterrupt() + + with pytest.raises(KeyboardInterrupt, match="peer map STORE cut"): + agent.add_remote_agent(b"peer-metadata", peer_name="decode-agent") + + peer = endpoint.open_peers[0] + assert agent._peers["decode-agent"] is peer + assert agent._peer_wire_metadata == {"decode-agent": b"peer-metadata"} + assert agent._source_aliases["remote-id"] == "decode-agent" + assert agent._source_aliases[("remote-id", "remote-incarnation")] == "decode-agent" + assert agent._construction_dirty is False + assert ( + agent.add_remote_agent(b"peer-metadata", peer_name="decode-agent") + == "decode-agent" + ) + assert endpoint.open_peers == (peer,) + + +def test_cached_plan_postpublication_cut_retries_without_duplicate_core_child(): + class StoreThenInterrupt(dict): + def __setitem__(self, key, value): + super().__setitem__(key, value) + raise KeyboardInterrupt("plan cache STORE cut") + + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") + remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") + agent._plan_cache = StoreThenInterrupt() + + with pytest.raises(KeyboardInterrupt, match="plan cache STORE cut"): + agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) + + plan = endpoint.open_plans[0] + cached = next(iter(agent._plan_cache.values())) + assert cached.plan is plan + assert agent._construction_dirty is False + handle = agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) + assert handle.plan is plan + assert endpoint.open_plans == (plan,) + + +def test_tracking_store_cut_cannot_hide_handle_from_worker_batch(): + class StoreThenInterrupt(dict): + armed = True + + def __setitem__(self, key, value): + super().__setitem__(key, value) + if self.armed: + self.armed = False + raise KeyboardInterrupt("work store cut") + + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + agent._works = StoreThenInterrupt() + handles = [] + agent.begin_handle_batch(handles) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + + with pytest.raises(KeyboardInterrupt, match="work store cut"): + agent.initialize_xfer("WRITE", local, remote, "decode-agent") + + assert handles == [] + assert agent._works == {} + assert endpoint.plans[-1].closed is True + agent.end_handle_batch(handles) + + +def test_begin_batch_store_cut_rolls_back_both_publications(): + class StoreThenInterruptLocal: + def __init__(self): + object.__setattr__(self, "armed", True) + + def __setattr__(self, name, value): + object.__setattr__(self, name, value) + if name == "value" and self.armed: + object.__setattr__(self, "armed", False) + raise KeyboardInterrupt("batch TLS store cut") + + agent, _, _ = _make_agent() + handles = [] + agent._batch_handles = StoreThenInterruptLocal() + + with pytest.raises(KeyboardInterrupt, match="batch TLS store cut"): + agent.begin_handle_batch(handles) + + assert agent._active_handle_batches == {} + assert getattr(agent._batch_handles, "value", None) is None + + +def test_post_tracking_cleanup_cut_is_retained_and_reaped(): + agent, endpoint, _ = _make_agent() + _register_and_add_peer(agent) + handles = [] + agent.begin_handle_batch(handles) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + original_prepare = endpoint.prepare + original_track = agent._track_work + + def prepare_with_one_close_cut(**kwargs): + plan = original_prepare(**kwargs) + plan.close_errors.append(RuntimeError("plan close cut")) + return plan + + def track_then_interrupt(handle): + original_track(handle) + raise KeyboardInterrupt("tracking return cut") + + with ( + patch.object(endpoint, "prepare", side_effect=prepare_with_one_close_cut), + patch.object(agent, "_track_work", side_effect=track_then_interrupt), + patch.object(agent, "_start_retained_work_sweeper"), + ): + with pytest.raises(KeyboardInterrupt, match="tracking return cut"): + agent.initialize_xfer("WRITE", local, remote, "decode-agent") + + handle = handles[0] + assert agent._works[id(handle)] is handle + assert id(handle) in agent._retained_work_ids + assert handle.owned_plan.closed is False + agent.end_handle_batch(handles) + with agent._lock: + agent._reap_retained_works() + + assert handle.phase == "released" + assert handle.owned_plan.closed is True + assert agent._retained_work_ids == set() + + +def test_whole_batch_retention_reaps_without_per_handle_journal(): + agent, _, _ = _make_agent() + _register_and_add_peer(agent) + handles = [] + agent.begin_handle_batch(handles) + local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") + remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") + handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") + assert handles == [handle] + assert agent.transfer(handle) == "PROC" + handle.work.cancel_completes = False + + with patch.object(agent, "_start_retained_work_sweeper"): + agent.retain_handle_batch(handles) + agent.end_handle_batch(handles) + assert agent._retained_handle_batches[id(handles)] is handles + assert id(handle) not in agent._retained_work_ids + handle.work.done = True + handle.work.state = _WorkState.COMPLETED + with agent._lock: + agent._reap_retained_works() + + assert handle.phase == "released" + assert agent._retained_handle_batches == {} + + +def test_retiring_peer_cleanup_retries_before_same_metadata_reimport(): + agent, endpoint, _ = _make_agent() + agent.add_remote_agent(b"peer-metadata", peer_name="decode-agent") + old_peer = endpoint.peers[-1] + old_peer.close_errors.append(KeyboardInterrupt("peer close cut")) + + with pytest.raises(KeyboardInterrupt, match="peer close cut"): + agent.rollback_remote_agent("decode-agent", b"peer-metadata") + + assert "decode-agent" in agent._retiring_peers + assert ( + agent.add_remote_agent(b"peer-metadata", peer_name="decode-agent") + == "decode-agent" + ) + assert old_peer.closed is True + assert endpoint.peers[-1] is not old_peer + assert "decode-agent" not in agent._retiring_peers diff --git a/test/registered/unit/managers/test_io_struct.py b/test/registered/unit/managers/test_io_struct.py index c35ba141acfd..a756381a2d9a 100644 --- a/test/registered/unit/managers/test_io_struct.py +++ b/test/registered/unit/managers/test_io_struct.py @@ -43,6 +43,34 @@ class TestTokenizedReqInputMsgpack(unittest.TestCase): + def test_transfer_generation_is_append_only_and_round_trips(self): + generation = "11" * 16 + self.assertEqual( + TokenizedGenerateReqInput.__struct_fields__[-1], "transfer_generation" + ) + request = TokenizedGenerateReqInput( + input_text="", + input_ids=array("q", [1, 2]), + input_embeds=None, + mm_inputs=None, + token_type_ids=None, + sampling_params=SamplingParams(), + return_logprob=False, + logprob_start_len=0, + top_logprobs_num=0, + token_ids_logprob=None, + stream=False, + routing_key="existing-wire-field", + cache_salt="existing-tail-field", + transfer_generation=generation, + ) + + decoded = msgpack_decode(msgpack_encode(request)) + + self.assertEqual(decoded.routing_key, "existing-wire-field") + self.assertEqual(decoded.cache_salt, "existing-tail-field") + self.assertEqual(decoded.transfer_generation, generation) + def test_rust_tokenized_generate_schema_stays_in_lockstep(self): """Compare the Rust wire declaration with the imported Python schema.""" rust_path = ( From db82eef2025235f539f1625acf465907ad4d9f90 Mon Sep 17 00:00:00 2001 From: Samuel Nordmann Date: Thu, 17 Sep 2026 16:12:09 +0200 Subject: [PATCH 2/2] [PD] Reduce Core transfer experiment to an opt-in hook Restore shared lifecycle, wire and staging code to the review base. Wrap the existing NIXL agent only when explicitly enabled; retain the native default. Assisted-by: OpenAI Codex Signed-off-by: Samuel Nordmann --- python/sglang/srt/disaggregation/base/conn.py | 5 +- .../sglang/srt/disaggregation/common/conn.py | 398 +-- .../disaggregation/common/staging_handler.py | 221 +- .../sglang/srt/disaggregation/common/utils.py | 2 +- python/sglang/srt/disaggregation/decode.py | 221 +- .../srt/disaggregation/mooncake/conn.py | 608 ++--- python/sglang/srt/disaggregation/nixl/conn.py | 1100 ++------ .../srt/disaggregation/nixl/torch_transfer.py | 2169 ---------------- .../disaggregation/nixl/torch_transfer_mvp.md | 13 + python/sglang/srt/disaggregation/prefill.py | 8 - python/sglang/srt/environ.py | 11 +- python/sglang/srt/managers/io_struct.py | 4 - python/sglang/srt/managers/schedule_batch.py | 4 - python/sglang/srt/managers/scheduler.py | 18 +- .../sglang/srt/managers/tokenizer_manager.py | 4 - rust/sglang-server/src/message/io_struct.rs | 54 +- rust/sglang-server/src/message/request.rs | 30 - .../test_deferred_decode_kv_release.py | 366 +-- .../test_disaggregation_wire.py | 17 +- .../test_mooncake_generation_control.py | 192 -- .../disaggregation/test_nixl_backend_basic.py | 846 +----- .../test_nixl_deferred_kv_release.py | 276 +- .../test_nixl_sender_failure_cleanup.py | 12 +- .../test_nixl_torch_transfer.py | 2257 ----------------- .../unit/managers/test_io_struct.py | 28 - 25 files changed, 751 insertions(+), 8113 deletions(-) delete mode 100644 python/sglang/srt/disaggregation/nixl/torch_transfer.py create mode 100644 python/sglang/srt/disaggregation/nixl/torch_transfer_mvp.md delete mode 100644 test/registered/unit/disaggregation/test_mooncake_generation_control.py delete mode 100644 test/registered/unit/disaggregation/test_nixl_torch_transfer.py diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index e97af5a9d1ac..358ff95c4304 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -104,10 +104,7 @@ class KVPoll: class BaseKVManager(ABC): """Base class for managing transfer states""" - # Decode destination memory is always quarantined after a transfer abort - # until every possible remote writer proves quiescence. Concrete managers - # must not make this safety property optional. - enable_deferred_decode_kv_release: bool = True + enable_deferred_decode_kv_release: bool = False @abstractmethod def __init__( diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 5c2439dbb52f..39763024a31b 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -3,13 +3,11 @@ import asyncio import concurrent.futures import dataclasses -import hashlib import logging -import secrets import threading import time from collections import defaultdict -from typing import Dict, FrozenSet, List, Optional, Set, Tuple, Union +from typing import Dict, List, Optional, Set, Tuple, Union import numpy as np import numpy.typing as npt @@ -52,57 +50,6 @@ logger = logging.getLogger(__name__) -_TRANSFER_GENERATION_NBYTES = 16 - - -def validate_transfer_generation(value: object) -> str: - """Return one canonical, wire-safe 128-bit decode-room lease identity.""" - - if type(value) is not str or len(value) != 2 * _TRANSFER_GENERATION_NBYTES: - raise ValueError("transfer generation must be a 32-character string") - if value.lower() != value: - raise ValueError("transfer generation must use canonical lowercase hex") - try: - decoded = bytes.fromhex(value) - except ValueError as exc: - raise ValueError( - "transfer generation must use canonical lowercase hex" - ) from exc - if len(decoded) != _TRANSFER_GENERATION_NBYTES: - raise ValueError("transfer generation must encode exactly 128 bits") - return value - - -def next_transfer_generation(value: object) -> str: - """Derive the next shared identity for a true-retraction rebootstrap. - - The initial value is minted randomly at ingress. Chaining a domain-separated - 128-bit BLAKE2 digest lets every decode TP rank independently derive the - same fresh lease without reusing the retired generation or coordinating a - new random value through the data plane. - """ - - generation = validate_transfer_generation(value) - return hashlib.blake2b( - bytes.fromhex(generation), - digest_size=_TRANSFER_GENERATION_NBYTES, - person=b"sgl-pd-reboot-v1", - ).hexdigest() - - -@dataclasses.dataclass(frozen=True) -class _AbortAckTarget: - decode_ip: str - decode_port: int - - -@dataclasses.dataclass -class _DeferredAbortTracker: - """Exact, immutable writer set and the ACKs adopted for one room lease.""" - - expected_prefill_ranks: FrozenSet[int] - acked_prefill_ranks: Set[int] = dataclasses.field(default_factory=set) - # 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 @@ -217,10 +164,9 @@ def __init__( self.is_hybrid_mla_backend = getattr(args, "is_hybrid_mla_backend", False) self.disaggregation_mode = disaggregation_mode self.server_args = server_args - # Compatibility/readback only. Correctness paths no longer branch on - # this value, so even an in-process mutation cannot bypass quarantine. - self.enable_deferred_decode_kv_release = True - self._abort_state_lock = threading.Lock() + self.enable_deferred_decode_kv_release = ( + envs.SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE.get() + ) self._dcp_pack_buffers = None # for p/d multi node infer self.bootstrap_host = get_serving().host @@ -297,12 +243,7 @@ def __init__( self.transfer_infos = {} # Deferred KV release: aborted room -> (decode_ip, decode_port); # ack held until the transfer drains. - self._deferred_ack_targets: Dict[Tuple[int, str], _AbortAckTarget] = {} - self._active_transfer_generations: Dict[int, str] = {} - # Without ACK-of-ACK, eviction could discard the only proof needed - # to replay a lost-return ACK. These cold-path tombstones therefore - # remain for the process lifetime. - self._abort_ack_receipts: Set[Tuple[int, str]] = set() + self._deferred_ack_targets: Dict[int, Tuple[str, int]] = {} self.req_to_decode_prefix_len: Dict[int, int] = {} self.decode_kv_args_table = {} self.pp_group = get_pp_group() @@ -324,9 +265,7 @@ def __init__( # 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[ - Tuple[int, str], _DeferredAbortTracker - ] = {} + self._deferred_abort_ack_tracker: Dict[int, Set[int]] = {} # Heartbeat interval should be at least 2 seconds self.heartbeat_interval = max( envs.SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL.get(), 2.0 @@ -444,145 +383,27 @@ def record_failure(self, bootstrap_room: int, failure_reason: str): with self.failure_lock: self.failure_records[bootstrap_room] = failure_reason - def register_deferred_abort_room( - self, - bootstrap_room: int, - transfer_generation: str, - expected_prefill_ranks: FrozenSet[int], - ) -> None: - """Publish the exact abort intent before any ABORT send can complete.""" - - key = (int(bootstrap_room), validate_transfer_generation(transfer_generation)) - expected = frozenset(expected_prefill_ranks) - if not expected or any(type(rank) is not int or rank < 0 for rank in expected): - raise ValueError( - "expected prefill ranks must be non-empty non-negative ints" - ) - with self._abort_state_lock: - tracker = self._deferred_abort_ack_tracker.get(key) - if tracker is None: - self._deferred_abort_ack_tracker[key] = _DeferredAbortTracker(expected) - elif tracker.expected_prefill_ranks != expected: - raise RuntimeError("abort intent cannot change its expected writer set") - - def note_abort_ack( - self, bootstrap_room: int, transfer_generation: str, prefill_rank: int - ) -> bool: - """Adopt an ACK only for the exact currently quarantined room lease.""" - - try: - generation = validate_transfer_generation(transfer_generation) - except ValueError: - return False - if type(prefill_rank) is not int or prefill_rank < 0: - return False - key = (int(bootstrap_room), generation) - with self._abort_state_lock: - tracker = self._deferred_abort_ack_tracker.get(key) - if tracker is None or prefill_rank not in tracker.expected_prefill_ranks: - return False - tracker.acked_prefill_ranks.add(prefill_rank) - return True - - def is_abort_release_safe( - self, - bootstrap_room: int, - transfer_generation: str, - ) -> bool: - """True only after every possible writer ACKs this exact lease.""" - - key = (int(bootstrap_room), validate_transfer_generation(transfer_generation)) - with self._abort_state_lock: - tracker = self._deferred_abort_ack_tracker.get(key) - return bool( - tracker is not None - and tracker.acked_prefill_ranks == tracker.expected_prefill_ranks - ) - - def clear_deferred_abort_state( - self, bootstrap_room: int, transfer_generation: str - ) -> None: - key = (int(bootstrap_room), validate_transfer_generation(transfer_generation)) - with self._abort_state_lock: - self._deferred_abort_ack_tracker.pop(key, None) - - def register_transfer_generation( - self, bootstrap_room: int, transfer_generation: str - ) -> bool: - """Atomically bind fresh metadata, rejecting retired or active ABA.""" - - room = int(bootstrap_room) - generation = validate_transfer_generation(transfer_generation) - with self._abort_state_lock: - if (room, generation) in self._abort_ack_receipts: - return False - current = self._active_transfer_generations.get(room) - if current is not None and current != generation: - return False - self._active_transfer_generations[room] = generation - return True + 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 active_transfer_generation(self, bootstrap_room: int) -> Optional[str]: - with self._abort_state_lock: - return self._active_transfer_generations.get(int(bootstrap_room)) + 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 _remember_abort_receipt_unlocked(self, key: Tuple[int, str]) -> None: - self._abort_ack_receipts.add(key) - - def mark_transfer_quiescent( - self, bootstrap_room: int, transfer_generation: Optional[str] = None - ) -> bool: - """Journal exact remote-writer quiescence before room teardown. - - The writer-count test and receipt publication share one critical - section with writer admission. A caller can therefore never retire a - generation between a sibling worker's increment and its DMA. - """ - - room = int(bootstrap_room) - with self._abort_state_lock: - if self._staging_outstanding.get(room, 0) > 0: - return False - generation = transfer_generation or self._active_transfer_generations.get( - room - ) - if generation is None: - return False - generation = validate_transfer_generation(generation) - key = (room, generation) - self._remember_abort_receipt_unlocked(key) - if self._active_transfer_generations.get(room) == generation: - self._active_transfer_generations.pop(room, None) - self._maybe_ack_drained_abort(room, generation) - return True - - def _count_transfer_writer(self, room: int) -> int: - """Admit one dequeued writer into the exact shared drain count.""" - - room = int(room) - with self._abort_state_lock: - count = self._staging_outstanding.get(room, 0) + 1 - self._staging_outstanding[room] = count - return count - - def _uncount_transfer_writer(self, room: int) -> int: - """Retire exactly one writer, never a concurrently active sibling.""" - - room = int(room) - with self._abort_state_lock: - count = self._staging_outstanding.get(room, 0) - if count <= 0: - raise RuntimeError(f"transfer-writer count underflow for room {room}") - remaining = count - 1 - if remaining: - self._staging_outstanding[room] = remaining - else: - self._staging_outstanding.pop(room, None) - return remaining + 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 + ) - def _transfer_writer_count(self, room: int) -> int: - with self._abort_state_lock: - return self._staging_outstanding.get(int(room), 0) + def clear_deferred_abort_state(self, bootstrap_room: int) -> None: + self._deferred_abort_ack_tracker.pop(bootstrap_room, None) def _prefill_unique_rank(self) -> int: """Stable per-sender id, matching what the transfer worker syncs on Success.""" @@ -592,14 +413,8 @@ def _prefill_unique_rank(self) -> int: + self.attn_cp_rank ) - def _send_abort_ack( - self, - decode_ip: str, - decode_port: int, - room: int, - transfer_generation: str, - ) -> bool: - """Send an idempotent generation-keyed ACK; report normal return only.""" + def _send_abort_ack(self, decode_ip: str, decode_port: int, room: int) -> None: + """Best-effort ack that this rank's transfer for an aborted room drained.""" try: na = NetworkAddress(decode_ip, decode_port) self._send_multipart_locked( @@ -608,82 +423,28 @@ def _send_abort_ack( b"ABORT_ACK", str(room).encode("ascii"), str(self._prefill_unique_rank()).encode("ascii"), - transfer_generation.encode("ascii"), ], is_ipv6=na.is_ipv6, ) - return True except Exception as e: - logger.debug( - "Failed to send drained ABORT_ACK for room %s generation %s: %s", - room, - transfer_generation, - e, - ) - return False + logger.debug(f"Failed to send drained ABORT_ACK for room {room}: {e}") - def _maybe_ack_drained_abort( - self, room: int, transfer_generation: Optional[str] = None - ) -> None: - """Replay a durable ACK only after exact-generation quiescence.""" - with self._abort_state_lock: - if self._staging_outstanding.get(int(room), 0) > 0: - return - if transfer_generation is None: - keys = tuple( - key for key in self._deferred_ack_targets if key[0] == int(room) - ) - else: - generation = validate_transfer_generation(transfer_generation) - keys = ((int(room), generation),) - sends = [] - for key in keys: - target = self._deferred_ack_targets.get(key) - if target is None: - continue - if key not in self._abort_ack_receipts: - active = self._active_transfer_generations.get(int(room)) - status = self.request_status.get(int(room)) - if active != key[1] or status != KVPoll.Failed: - continue - self._remember_abort_receipt_unlocked(key) - self._active_transfer_generations.pop(int(room), None) - sends.append((key, target)) - for key, target in sends: - if self._send_abort_ack( - target.decode_ip, - target.decode_port, - key[0], - key[1], - ): - with self._abort_state_lock: - if self._deferred_ack_targets.get(key) == target: - self._deferred_ack_targets.pop(key, None) + 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 not None: + self._send_abort_ack(target[0], target[1], room) def register_deferred_ack_target( - self, - room: int, - transfer_generation: str, - decode_ip: str, - decode_port: int, - ) -> bool: - """Journal a target only for an active or already-quiescent exact lease.""" - - room = int(room) - generation = validate_transfer_generation(transfer_generation) - key = (room, generation) - target = _AbortAckTarget(str(decode_ip), int(decode_port)) - with self._abort_state_lock: - if ( - self._active_transfer_generations.get(room) != generation - and key not in self._abort_ack_receipts - ): - return False - current_target = self._deferred_ack_targets.get(key) - if current_target is not None and current_target != target: - return False - self._deferred_ack_targets[key] = target - return True + self, room: int, decode_ip: str, decode_port: int + ) -> 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.""" + self._deferred_ack_targets[room] = (decode_ip, decode_port) def get_kv_replica_factor(self) -> int: if self._kv_replica_factor is None: @@ -1580,8 +1341,10 @@ def clear(self) -> None: 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) - # A generation-keyed ACK target is a durable ownership journal. Only a - # normally returned ACK send may remove it; sender clear is not proof. + if hasattr(self.kv_mgr, "_deferred_ack_targets"): + # Drop a held ack target if the room concluded without draining + # (e.g. aborted before any chunk enqueued); else it leaks on prefill. + self.kv_mgr._deferred_ack_targets.pop(self.bootstrap_room, None) def abort(self): self.kv_mgr.record_failure( @@ -1604,18 +1367,8 @@ def __init__( mgr: CommonKVManager, bootstrap_addr: str, bootstrap_room: Optional[int] = None, - transfer_generation: Optional[str] = None, ): self.bootstrap_room = bootstrap_room - # Current request plumbing supplies one value before scheduler fanout, - # so every decode rank targeting the same prefill manager uses one - # lease identity. The fallback preserves direct legacy construction. - generation = ( - secrets.token_hex(_TRANSFER_GENERATION_NBYTES) - if transfer_generation is None - else validate_transfer_generation(transfer_generation) - ) - self._transfer_generation = generation self.bootstrap_addr = bootstrap_addr self.kv_mgr = mgr self.conclude_state: Optional[KVPoll] = None @@ -1626,12 +1379,6 @@ def __init__( self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping) - @property - def transfer_generation(self) -> str: - """Immutable generation shared by all ranks of this room lease.""" - - return self._transfer_generation - def init(self, prefill_dp_rank: int): if self.bootstrap_addr not in self.kv_mgr.prefill_info_table: self.kv_mgr.record_failure( @@ -1690,15 +1437,6 @@ def _setup_bootstrap_infos(self): target_pp_rank, ) if bootstrap_info is not None: - bootstrap_info["prefill_unique_rank"] = ( - target_tp_rank - * ( - self.prefill_info.pp_size - * self.prefill_info.attn_cp_size - ) - + target_pp_rank * self.prefill_info.attn_cp_size - + target_cp_rank - ) if self.kv_mgr.is_mla_backend: # For MLA: target_tp_rank is the selected real rank, others are dummy ranks bootstrap_info["is_dummy"] = not bool( @@ -1870,9 +1608,13 @@ def _check_waiting_timeout(self) -> Optional[KVPoll]: ) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) self.invalidate_cached_bootstrap_infos() - if hasattr(self, "bootstrap_infos") and self.bootstrap_infos is not None: - self.arm_abort_intent() + if ( + not self.abort_notified + and hasattr(self, "bootstrap_infos") + and self.bootstrap_infos is not None + ): self._send_abort_notification() + self.abort_notified = True return KVPoll.Failed def clear(self) -> None: @@ -1881,52 +1623,19 @@ def clear(self) -> None: self.kv_mgr.prefill_response_tracker.pop(self.bootstrap_room, None) def abort(self): - self.arm_abort_intent() self.kv_mgr.record_failure( self.bootstrap_room, "Aborted by AbortReq.", ) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) self.conclude_state = KVPoll.Failed - if hasattr(self, "bootstrap_infos") and self.bootstrap_infos is not None: - self._send_abort_notification() - - def arm_abort_intent(self) -> None: - """Publish quarantine ownership before the first fallible wire send.""" - - if self.abort_notified: - return - bootstrap_infos = getattr(self, "bootstrap_infos", None) - if not bootstrap_infos: - # Without rank discovery there is no authenticated writer set. Any - # deferred owner therefore remains quarantined and cannot become - # release-safe from a fabricated empty set. - logger.error( - "Cannot arm ABORT for room %s generation %s before exact " - "prefill-rank discovery; decode pages will remain quarantined", - self.bootstrap_room, - self.transfer_generation, - ) - return - expected_prefill_ranks = frozenset( - bootstrap_info["prefill_unique_rank"] for bootstrap_info in bootstrap_infos - ) - self.kv_mgr.register_deferred_abort_room( - self.bootstrap_room, - self.transfer_generation, - expected_prefill_ranks, - ) - self.abort_notified = True - - def retry_abort_notification(self) -> None: - """Replay an armed ABORT; duplicate delivery is generation-idempotent.""" - if ( - self.abort_notified + not self.abort_notified and hasattr(self, "bootstrap_infos") and self.bootstrap_infos is not None ): self._send_abort_notification() + self.abort_notified = True def _send_abort_notification(self): for bootstrap_info in self.bootstrap_infos: @@ -1940,7 +1649,6 @@ def _send_abort_notification(self): str(self.bootstrap_room).encode("ascii"), self.kv_mgr.local_ip.encode("ascii"), str(self.kv_mgr.rank_port).encode("ascii"), - self.transfer_generation.encode("ascii"), ] ) logger.debug( diff --git a/python/sglang/srt/disaggregation/common/staging_handler.py b/python/sglang/srt/disaggregation/common/staging_handler.py index f68f50a347a1..a9ae31b850d5 100644 --- a/python/sglang/srt/disaggregation/common/staging_handler.py +++ b/python/sglang/srt/disaggregation/common/staging_handler.py @@ -17,7 +17,6 @@ import torch -from sglang.srt.disaggregation.common.conn import validate_transfer_generation from sglang.srt.runtime_context import ( get_schedule, ) @@ -27,10 +26,6 @@ # Bounded wait for a watermark advance before re-enqueueing a deferred staging # chunk, so the re-enqueue retry does not busy-spin a core. STAGING_WATERMARK_WAIT_S = 0.001 -STAGING_REQ_WIRE_TAG = b"STAGING_REQ_V2" -STAGING_RSP_WIRE_TAG = b"STAGING_RSP_V2" -_MAX_STAGING_INT_DIGITS = 20 -_MAX_STAGING_SESSION_BYTES = 4096 if TYPE_CHECKING: from sglang.srt.disaggregation.decode import DecodeRequest @@ -104,10 +99,8 @@ def __init__( # before unregister runs, but release_room still needs it. self._room_to_receiver: dict = {} self._wm_subscribers: dict = {} - # room -> chunk_idx -> {exact_writer_id: (page_start, num_pages)} fan-in; - # handler-owned so room teardown can purge it. A map is required here: - # duplicate completion traffic from one rank must not impersonate the - # other writers needed to authorize scatter. + # room -> chunk_idx -> [(page_start, num_pages, writer_id)] fan-in + # arrivals; handler-owned so room teardown can purge them. self._writer_counts: dict = {} def register_wm_subscriber(self, receiver, session_id: str) -> None: @@ -314,37 +307,11 @@ def handle_chunk_arrived( ) return False room_counts = self._writer_counts.setdefault(room, {}) - arrivals = room_counts.setdefault(chunk_idx, {}) - arrival = (page_start, num_pages) - previous = arrivals.get(writer_id) - if previous is not None: - if previous != arrival: - logger.error( - "Conflicting duplicate staging completion room=%s chunk=%s " - "writer=%s prior=%s current=%s", - room, - chunk_idx, - writer_id, - previous, - arrival, - ) - return False - arrivals[writer_id] = arrival + arrivals = room_counts.setdefault(chunk_idx, []) + arrivals.append((page_start, num_pages, writer_id)) num_writers = self.num_writers_for(receiver) if len(arrivals) >= num_writers: - layouts = set(arrivals.values()) - if len(layouts) != 1: - logger.error( - "Conflicting staging completion layouts room=%s chunk=%s: %s", - room, - chunk_idx, - layouts, - ) - return False - scatter_page_start, scatter_num_pages = next(iter(layouts)) - self.submit_chunk_scatter( - room, chunk_idx, scatter_page_start, scatter_num_pages - ) + self.submit_chunk_scatter(room, chunk_idx, page_start, num_pages) del room_counts[chunk_idx] return True return False @@ -535,99 +502,27 @@ def handle_watermark_msg(staging_ctx, msg_parts) -> None: staging_ctx.watermark_cv.notify_all() -def _parse_staging_int(frame, field: str, *, minimum: int = 0) -> int: - if type(frame) is not bytes: - raise ValueError(f"{field} must be exact bytes") - value = frame.decode("ascii") - digits = value[1:] if value.startswith("-") else value - if ( - not digits - or len(digits) > _MAX_STAGING_INT_DIGITS - or not digits.isdecimal() - or (len(digits) > 1 and digits[0] == "0") - or value == "-0" - ): - raise ValueError(f"{field} must be canonical bounded decimal") - parsed = int(value) - if parsed < minimum: - raise ValueError(f"{field} must be at least {minimum}") - return parsed - - -def _parse_staging_session(frame) -> str: - if type(frame) is not bytes or not frame or len(frame) > _MAX_STAGING_SESSION_BYTES: - raise ValueError("staging session must be non-empty bounded bytes") - return frame.decode("ascii") - - -def _parse_staging_req(msg): - if len(msg) not in (6, 7) or msg[0] != STAGING_REQ_WIRE_TAG: - raise ValueError("unsupported STAGING_REQ wire layout") - generation = validate_transfer_generation(msg[1].decode("ascii")) - room = _parse_staging_int(msg[2], "staging room") - chunk_idx = _parse_staging_int(msg[3], "staging chunk") - chunk_num_pages = _parse_staging_int(msg[4], "staging chunk page count", minimum=1) - session_id = _parse_staging_session(msg[5]) - requester_pp_rank = ( - _parse_staging_int(msg[6], "requester PP rank") if len(msg) == 7 else None - ) - return ( - generation, - room, - chunk_idx, - chunk_num_pages, - session_id, - requester_pp_rank, - ) - - -def handle_staging_rsp(msg_parts, transfer_infos: dict) -> bool: - """Adopt an allocation response only for the exact active room lease.""" - try: - if len(msg_parts) != 8 or msg_parts[0] != STAGING_RSP_WIRE_TAG: - raise ValueError("unsupported STAGING_RSP wire layout") - generation = validate_transfer_generation(msg_parts[1].decode("ascii")) - stg_room = _parse_staging_int(msg_parts[2], "staging room") - stg_chunk_idx = _parse_staging_int(msg_parts[3], "staging chunk") - stg_offset = _parse_staging_int(msg_parts[4], "staging offset", minimum=-2) - stg_round = _parse_staging_int(msg_parts[5], "staging round") - stg_end = _parse_staging_int(msg_parts[6], "staging end", minimum=-1) - stg_session = _parse_staging_session(msg_parts[7]) - except (AttributeError, IndexError, TypeError, UnicodeError, ValueError) as error: - logger.warning("Dropping malformed or legacy STAGING_RSP: %s", error) - return False - +def handle_staging_rsp(msg_parts, transfer_infos: dict) -> None: + """Process a STAGING_RSP message and update transfer info with allocation.""" + stg_room = int(msg_parts[1].decode("ascii")) + stg_chunk_idx = int(msg_parts[2].decode("ascii")) + stg_offset = int(msg_parts[3].decode("ascii")) + stg_round = int(msg_parts[4].decode("ascii")) + stg_end = int(msg_parts[5].decode("ascii")) + stg_session = msg_parts[6].decode("ascii") room_infos = transfer_infos.get(stg_room, {}) tinfo = room_infos.get(stg_session) - if tinfo is None: + if tinfo is not None: + if tinfo.staging is None: + tinfo.staging = StagingTransferInfo() + tinfo.staging.set_chunk(stg_chunk_idx, stg_offset, stg_round, stg_end) + else: logger.warning( "STAGING_RSP RECV but tinfo=None room=%s chunk=%d session=%s", stg_room, stg_chunk_idx, stg_session, ) - return False - try: - active_generation = validate_transfer_generation(tinfo.transfer_generation) - except (AttributeError, TypeError, ValueError): - logger.warning( - "Dropping STAGING_RSP for room=%s with no authenticated active lease", - stg_room, - ) - return False - if generation != active_generation: - logger.warning( - "Dropping stale STAGING_RSP room=%s response_generation=%s " - "active_generation=%s", - stg_room, - generation, - active_generation, - ) - return False - if tinfo.staging is None: - tinfo.staging = StagingTransferInfo() - tinfo.staging.set_chunk(stg_chunk_idx, stg_offset, stg_round, stg_end) - return True # ====================================================================== @@ -873,9 +768,7 @@ def handle_staging_req( kv_buffer_tensors, room_receivers: dict, room_bootstrap: dict, - *, - parsed_request=None, -) -> bool: +): """Allocate staging for a chunk on-demand and send STAGING_RSP to prefill. Deduplicates: multiple prefill TP ranks requesting the same (room, chunk_idx) @@ -883,18 +776,11 @@ def handle_staging_req( """ from sglang.srt.disaggregation.common.staging_buffer import StagingAllocator - try: - ( - generation, - room, - chunk_idx, - chunk_num_pages, - session_id, - requester_pp_rank, - ) = parsed_request if parsed_request is not None else _parse_staging_req(msg) - except (AttributeError, IndexError, TypeError, UnicodeError, ValueError) as error: - logger.warning("Dropping malformed or legacy STAGING_REQ: %s", error) - return False + room = int(msg[1].decode("ascii")) + chunk_idx = int(msg[2].decode("ascii")) + chunk_num_pages = int(msg[3].decode("ascii")) + session_id = msg[4].decode("ascii") + requester_pp_rank = int(msg[5].decode("ascii")) if len(msg) > 5 else None if staging_allocator is None: logger.warning( @@ -902,7 +788,7 @@ def handle_staging_req( room, chunk_idx, ) - return False + return receiver = room_receivers.get(room) if receiver is None: @@ -912,23 +798,7 @@ def handle_staging_req( chunk_idx, session_id, ) - return False - try: - active_generation = validate_transfer_generation(receiver.transfer_generation) - except (AttributeError, TypeError, ValueError): - logger.warning( - "STAGING_REQ dropped: room=%s has no authenticated active lease", room - ) - return False - if generation != active_generation: - logger.warning( - "STAGING_REQ dropped: stale room=%s request_generation=%s " - "active_generation=%s", - room, - generation, - active_generation, - ) - return False + return infos = receiver.chunk_staging_infos if chunk_idx < len(infos) and infos[chunk_idx][0] >= 0: @@ -1000,8 +870,7 @@ def handle_staging_req( with lock: sock.send_multipart( [ - STAGING_RSP_WIRE_TAG, - generation.encode("ascii"), + b"STAGING_RSP", str(room).encode("ascii"), str(chunk_idx).encode("ascii"), str(offset).encode("ascii"), @@ -1012,7 +881,6 @@ def handle_staging_req( ) except Exception: pass - return True class StagingManagerMixin: @@ -1030,18 +898,8 @@ def _is_watermark_ready( return is_watermark_ready(self._staging_ctx, session_id, alloc_round, alloc_end) def _handle_staging_req(self, msg): - try: - parsed_request = _parse_staging_req(msg) - except ( - AttributeError, - IndexError, - TypeError, - UnicodeError, - ValueError, - ) as error: - logger.warning("Dropping malformed or legacy STAGING_REQ: %s", error) - return - generation, room, _, _, session_id, _ = parsed_request + room = int(msg[1].decode("ascii")) + session_id = msg[4].decode("ascii") handler = self._staging_handler assert handler is not None, ( "STAGING_REQ received before staging handler initialized" @@ -1053,17 +911,8 @@ def _handle_staging_req(self, msg): room, ) return - receiver = self._staging_ctx.room_receivers.get(room) - if receiver is None or receiver.transfer_generation != generation: - logger.warning( - "STAGING_REQ received for stale room lease room=%s generation=%s, " - "skipping", - room, - generation, - ) - return prefill_tp = decode_req.kv_receiver.prefill_info.attn_tp_size - accepted = handle_staging_req( + handle_staging_req( msg, self._staging_ctx.allocator, self.kv_args, @@ -1072,10 +921,10 @@ def _handle_staging_req(self, msg): getattr(self, "kv_buffer_tensors", None), self._staging_ctx.room_receivers, self._staging_ctx.room_bootstrap, - parsed_request=parsed_request, ) - if accepted: + receiver = self._staging_ctx.room_receivers.get(room) + if receiver is not None: handler.register_wm_subscriber(receiver, session_id) @@ -1116,8 +965,7 @@ def prefetch_staging_reqs( num_chunks = (total_pages + full_chunk_pages - 1) // full_chunk_pages for chunk_idx in range(num_chunks): - generation = validate_transfer_generation(tinfo.transfer_generation) - stg_key = (room, generation, chunk_idx, session_id) + stg_key = (room, chunk_idx, session_id) if stg_key in staging_requested: continue staging_requested.add(stg_key) @@ -1134,8 +982,7 @@ def prefetch_staging_reqs( sock.connect(ep) prefetch_sockets[ep] = sock request = [ - STAGING_REQ_WIRE_TAG, - generation.encode("ascii"), + b"STAGING_REQ", str(room).encode("ascii"), str(chunk_idx).encode("ascii"), str(chunk_pages).encode("ascii"), diff --git a/python/sglang/srt/disaggregation/common/utils.py b/python/sglang/srt/disaggregation/common/utils.py index d03713650738..1571c3d19372 100644 --- a/python/sglang/srt/disaggregation/common/utils.py +++ b/python/sglang/srt/disaggregation/common/utils.py @@ -32,7 +32,7 @@ class TransferKVChunk: # Set when the staging worker first counts this chunk toward the per-room # outstanding count; stays set across re-enqueue on a watermark defer. staging_counted: bool = False - # Early-send CUDA event to synchronize before reading source memory. + # Mori early-send: CUDA event to synchronize before RDMA (optional). wait_event: Optional[object] = None diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 67d17b7cc939..9d6d91966357 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -37,12 +37,7 @@ from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.disaggregation.base import KVPoll from sglang.srt.disaggregation.base.conn import StateType -from sglang.srt.disaggregation.common.conn import ( - CommonKVManager, - CommonKVReceiver, - next_transfer_generation, - validate_transfer_generation, -) +from sglang.srt.disaggregation.common.conn import CommonKVManager, CommonKVReceiver from sglang.srt.disaggregation.decode_hicache_mixin import ( DecodeHiCachePreallocMixin, DecodeHiCacheTransferMixin, @@ -636,35 +631,6 @@ def add( if self._check_if_req_exceed_kv_capacity(req): return - backend = ( - TransferBackend.FAKE if _is_fake_transfer(req) else self.transfer_backend - ) - if not is_retracted and backend in ( - TransferBackend.NIXL, - TransferBackend.MOONCAKE, - ): - try: - generation = validate_transfer_generation(req.transfer_generation) - except ValueError as error: - message = ( - "Missing or invalid internal transfer generation on a " - f"disaggregated request: {error}" - ) - logger.error(message) - prepare_abort( - req, - message, - status_code=HTTPStatus.INTERNAL_SERVER_ERROR, - ) - self.scheduler.output_streamer.stream_output([req], req.return_logprob) - return - if is_rebootstrap: - # The completed lease has a process-lifetime tombstone and its - # delayed notifications must never match this new transfer. - # All TP ranks start from the same ingress generation, so this - # deterministic chain advances them to the same next identity. - req.transfer_generation = next_transfer_generation(generation) - if is_retracted: req.retraction_mb_id = None self.retracted_queue.append(req) @@ -733,16 +699,10 @@ def _create_receiver_and_enqueue( ) kv_receiver_class = get_kv_class(backend, KVClassType.RECEIVER) - receiver_kwargs = ( - {"transfer_generation": req.transfer_generation} - if backend in (TransferBackend.NIXL, TransferBackend.MOONCAKE) - else {} - ) kv_receiver = kv_receiver_class( mgr=self.kv_manager, bootstrap_addr=_bootstrap_addr(req), bootstrap_room=req.bootstrap_room, - **receiver_kwargs, ) decode_req = DecodeRequest( @@ -2069,19 +2029,6 @@ def _generate_fake_prefill_handoff_output_id(req: Req) -> int: return int.from_bytes(digest, byteorder="little") % req.vocab_size -@dataclass -class _DeferredKVRelease: - decode_req: DecodeRequest - bootstrap_room: int - kv_receiver: CommonKVReceiver - next_diagnostic_at: float - metadata_buffer_index: int - transfer_generation: str - diagnostic_count: int = 0 - release_phase: str = "QUARANTINED" - release_error: Optional[Tuple[str, str]] = None - - class DecodeTransferQueue(DecodeHiCacheTransferMixin): """ Store the requests that is polling kv @@ -2106,13 +2053,15 @@ def __init__( self.spec_algorithm = scheduler.spec_algorithm self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() self.staging_handler = None - # Non-bypassable: a remote writer must prove quiescence before these - # pages or this request slot can be reused. - self.enable_deferred_kv_release = True + self.enable_deferred_kv_release = ( + envs.SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE.get() + ) self.deferred_kv_release_timeout = ( envs.SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE_TIMEOUT.get() ) - self._deferred_releases: List[_DeferredKVRelease] = [] + # Aborted-mid-transfer requests whose KV pages/slot are held until drained + # or timed out. Entries: (decode_req, deadline, metadata_idx, required_acks). + self._deferred_releases: List[Tuple[DecodeRequest, float, int, int]] = [] def add(self, decode_req: DecodeRequest) -> None: self.queue.append(decode_req) @@ -2384,9 +2333,13 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req ) if self.scheduler.enable_hisparse: self.scheduler.hisparse_coordinator.request_finished(decode_req.req) - if decode_req.kv_receiver.abort_notified: + if ( + self.enable_deferred_kv_release + and decode_req.kv_receiver.kv_mgr.enable_deferred_decode_kv_release + and decode_req.kv_receiver.abort_notified + ): # Decode-initiated abort: a prefill write may still target - # these pages, so hold them until the exact drain proof. + # these pages, so hold them until the drain ack or timeout. # (A prefill-initiated failure has already stopped writing -> # immediate release below.) self._defer_release(decode_req) @@ -2459,143 +2412,65 @@ def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req def _defer_release(self, decode_req: DecodeRequest) -> None: deadline = time.monotonic() + self.deferred_kv_release_timeout - receiver = decode_req.kv_receiver - receiver.arm_abort_intent() + # Require an ack from every notified prefill rank (dummy-proof). Snapshot + # now -- the receiver may be cleared by resolve time. + required_acks = len(decode_req.kv_receiver.bootstrap_infos) self._deferred_releases.append( - _DeferredKVRelease( - decode_req=decode_req, - bootstrap_room=decode_req.req.bootstrap_room, - kv_receiver=receiver, - next_diagnostic_at=deadline, - metadata_buffer_index=decode_req.metadata_buffer_index, - transfer_generation=receiver.transfer_generation, - ) + (decode_req, deadline, decode_req.metadata_buffer_index, required_acks) ) - def _do_release(self, held: _DeferredKVRelease) -> None: - """Run one fail-stop release transaction without replaying a free. - - Each phase is journaled before its fallible or non-idempotent action. If - an action commits and its return is lost, the journal blocks automatic - retry and prefers a diagnosable leak to a possible double-free. - """ - - if held.release_phase != "QUARANTINED": - raise RuntimeError( - f"unsafe deferred release replay from phase {held.release_phase}" - ) - decode_req = held.decode_req - receiver = held.kv_receiver - room = held.bootstrap_room - idx = held.metadata_buffer_index + def _do_release(self, decode_req: DecodeRequest, idx: int) -> None: + room = decode_req.req.bootstrap_room if self.enable_staging and self.staging_handler.is_staging_room(room): - held.release_phase = "STAGING_RELEASE_ENTERED" self.staging_handler.unregister_decode_req(room) # release pre-allocated kv cache, but don't insert into the tree since it's failed - held.release_phase = "KV_RELEASE_ENTERED" release_kv_cache(decode_req.req, self.tree_cache, is_insert=False) - held.release_phase = "METADATA_RESET_ENTERED" self.metadata_buffers.bootstrap_room[idx] = 0 - held.release_phase = "METADATA_INDEX_FREE_ENTERED" self.req_to_metadata_buffer_idx_allocator.free(idx) - held.release_phase = "CONTROL_CLEANUP_ENTERED" - receiver.kv_mgr.clear_deferred_abort_state(room, held.transfer_generation) - receiver.clear() + decode_req.kv_receiver.kv_mgr.clear_deferred_abort_state(room) + decode_req.kv_receiver.clear() decode_req.kv_receiver = None - held.release_phase = "DONE" def has_pending_deferred_releases(self) -> bool: return bool(self._deferred_releases) def resolve_deferred_releases(self) -> None: - """Release only on exact drain proof; timeout retries and diagnoses.""" + """Release held requests once every prefill rank acks the drain, or the + hold times out.""" if not self._deferred_releases: return now = time.monotonic() - for held in tuple(self._deferred_releases): - if held.release_phase == "DONE": - self._deferred_releases = [ - entry for entry in self._deferred_releases if entry is not held - ] - continue - if held.release_phase != "QUARANTINED": - if now >= held.next_diagnostic_at: - logger.error( - "Fail-stop KV release journal retains room %s generation " - "%s in phase %s after %s; automatic retry is unsafe", - held.bootstrap_room, - held.transfer_generation, - held.release_phase, - held.release_error, - ) - held.next_diagnostic_at = now + 300.0 - continue - decode_req = held.decode_req - receiver = held.kv_receiver - room = held.bootstrap_room - kv_mgr = receiver.kv_mgr - drained = kv_mgr.is_abort_release_safe(room, held.transfer_generation) - if drained: - try: - self._do_release(held) - except BaseException as error: - try: - error_text = str(error) - except BaseException: - error_text = "" - held.release_error = ( - f"{type(error).__module__}.{type(error).__qualname__}", - error_text[:512], - ) - held.next_diagnostic_at = now + 300.0 - logger.exception( - "Deferred KV release entered fail-stop phase %s for " - "room %s generation %s; owner remains quarantined", - held.release_phase, - room, - held.transfer_generation, - ) - if not isinstance(error, Exception): - raise - else: - self._deferred_releases = [ - entry for entry in self._deferred_releases if entry is not held - ] - continue - if now >= held.next_diagnostic_at: + still_held = [] + to_release = [] + for decode_req, deadline, idx, required_acks in self._deferred_releases: + room = decode_req.req.bootstrap_room + kv_mgr = decode_req.kv_receiver.kv_mgr + drained = kv_mgr.is_abort_release_safe(room, required_acks) + if not drained and now < deadline: + still_held.append((decode_req, deadline, idx, required_acks)) + else: + to_release.append((decode_req, idx, room, drained)) + # Commit the survivors before releasing so a _do_release exception can't + # leave a released entry in the list (double-free / None receiver on retry). + self._deferred_releases = still_held + for decode_req, idx, room, drained in to_release: + if not drained: logger.warning( - "Deferred KV quarantine for room %s generation %s has no " - "full drain proof after %.1fs; retaining pages and retrying " - "ABORT (diagnostic %d)", - room, - held.transfer_generation, - self.deferred_kv_release_timeout, - held.diagnostic_count + 1, + f"Deferred KV release for room {room} timed out after " + f"{self.deferred_kv_release_timeout}s without a full drain " + f"ack from prefill; releasing anyway." ) - try: - receiver.retry_abort_notification() - except Exception: - logger.exception( - "ABORT retry failed for quarantined room %s generation %s", - room, - held.transfer_generation, - ) - held.diagnostic_count += 1 - backoff = min( - max(self.deferred_kv_release_timeout, 1.0) - * (2 ** min(held.diagnostic_count, 5)), - 300.0, - ) - held.next_diagnostic_at = now + backoff + try: + self._do_release(decode_req, idx) + except Exception: + # Isolate a failed release so the rest still run; entry already dropped. + logger.exception(f"Deferred KV release failed for room {room}") def release_memory_occupation(self): """Clean up in-flight transfers before releasing GPU memory.""" - if self._deferred_releases: - raise RuntimeError( - "Cannot release decode GPU memory while remote-writer " - f"quarantine holds {len(self._deferred_releases)} request(s)" - ) self.queue.clear() + # Pool is being torn down; drop held entries without per-request release. + self._deferred_releases.clear() def resume_memory_occupation(self): """Queues are already cleared on release; new transfers can be accepted.""" diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index be1e4fe26cab..d8914b736d98 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -8,7 +8,7 @@ import threading import time from collections import defaultdict -from typing import FrozenSet, List, Optional, Set, Tuple, Union +from typing import List, Optional, Set, Tuple, Union import numpy as np import numpy.typing as npt @@ -22,12 +22,9 @@ CommonKVReceiver, CommonKVSender, KVTransferError, - validate_transfer_generation, ) from sglang.srt.disaggregation.common.staging_handler import ( STAGING_WATERMARK_WAIT_S, - STAGING_REQ_WIRE_TAG, - STAGING_RSP_WIRE_TAG, DecodeStagingContext, PrefillStagingContext, StagingManagerMixin, @@ -78,31 +75,6 @@ logger = logging.getLogger(__name__) -MOONCAKE_CHUNK_READY_WIRE_TAG = b"CHUNK_READY_V2" -MOONCAKE_AUX_DATA_WIRE_TAG = b"AUX_DATA_V2" -MOONCAKE_TRANSFER_STATUS_WIRE_TAG = b"TRANSFER_STATUS_V2" -_MAX_MOONCAKE_CONTROL_UINT_DIGITS = 20 - - -def _parse_mooncake_control_uint( - frame: bytes, field: str, *, positive: bool = False -) -> int: - if type(frame) is not bytes: - raise ValueError(f"{field} must be exact bytes") - value = frame.decode("ascii") - if ( - not value - or len(value) > _MAX_MOONCAKE_CONTROL_UINT_DIGITS - or not value.isdecimal() - or (len(value) > 1 and value[0] == "0") - ): - raise ValueError(f"{field} must be canonical bounded unsigned decimal") - parsed = int(value) - if positive and parsed == 0: - raise ValueError(f"{field} must be positive") - return parsed - - FAILED_SESSION_RECOVERIES = Counter( "sglang:failed_session_recoveries_total", "Number of mooncake_session_ids un-blacklisted via probe.", @@ -123,7 +95,6 @@ class TransferInfo: is_dummy: bool decode_prefix_len: Optional[int] = None dst_device_kv_indices: Optional[npt.NDArray[np.int32]] = None - transfer_generation: str = "" # Note: always put the optional staging field at the final (it will be set through 'STAGING_RSP' pkg when needed) staging: Optional[StagingTransferInfo] = None @@ -157,9 +128,6 @@ def from_zmq(cls, msg: List[bytes]): if len(msg) > 9 and msg[9] != b"" else None ), - transfer_generation=validate_transfer_generation( - msg[10].decode("ascii") if len(msg) > 10 else None - ), ) @@ -237,15 +205,8 @@ def from_zmq(cls, msg: List[bytes]): ) -@dataclasses.dataclass(frozen=True) -class _MooncakeReceiveLease: - generation: str - session_id: str - expected_prefill_ranks: Optional[FrozenSet[int]] = None - - class MooncakeKVManager(StagingManagerMixin, CommonKVManager): - AUX_DATA_HEADER = MOONCAKE_AUX_DATA_WIRE_TAG + AUX_DATA_HEADER = b"AUX_DATA" def __init__( self, @@ -325,7 +286,6 @@ def __init__( daemon=True, ).start() elif self.disaggregation_mode == DisaggregationMode.DECODE: - self._receive_transfer_leases: dict[int, _MooncakeReceiveLease] = {} self._staging_ctx = DecodeStagingContext() if self.enable_staging else None if self.enable_staging: self._init_staging_allocator() @@ -383,77 +343,6 @@ def register_staging_room_bootstrap(self, room, bootstrap_infos, receiver): self._staging_ctx.room_bootstrap[room] = bootstrap_infos self._staging_ctx.room_receivers[room] = receiver - def _register_receive_transfer_lease( - self, - room: int, - generation: str, - session_id: str, - expected_prefill_ranks: Optional[FrozenSet[int]] = None, - ) -> None: - """Bind immutable decode control identity before accepting wire events.""" - - room = int(room) - generation = validate_transfer_generation(generation) - if type(session_id) is not str or not session_id: - raise ValueError("Mooncake receive lease requires a non-empty session") - expected = None - if expected_prefill_ranks is not None: - expected = frozenset(expected_prefill_ranks) - if not expected or any( - type(rank) is not int or rank < 0 for rank in expected - ): - raise ValueError( - "Mooncake receive lease ranks must be non-empty non-negative ints" - ) - candidate = _MooncakeReceiveLease(generation, session_id, expected) - with self._abort_state_lock: - current = self._receive_transfer_leases.get(room) - if current is None: - self._receive_transfer_leases[room] = candidate - self.prefill_response_tracker[room] = set() - return - if current.generation != generation or current.session_id != session_id: - raise RuntimeError( - "Mooncake room already has a different receive lease" - ) - if current.expected_prefill_ranks is not None and current != candidate: - raise RuntimeError("Mooncake receive lease rank identity is immutable") - self._receive_transfer_leases[room] = candidate - - @staticmethod - def _receive_control_matches_unlocked( - lease: Optional[_MooncakeReceiveLease], - generation: str, - prefill_rank: int, - *, - session_id: Optional[str] = None, - ) -> bool: - return bool( - lease is not None - and lease.expected_prefill_ranks is not None - and lease.generation == generation - and prefill_rank in lease.expected_prefill_ranks - and (session_id is None or lease.session_id == session_id) - ) - - def _clear_receive_transfer_lease(self, room: int, generation: str) -> bool: - """Clear only this receiver's identity and room state, atomically.""" - - room = int(room) - generation = validate_transfer_generation(generation) - with self._abort_state_lock: - current = self._receive_transfer_leases.get(room) - # Absence is not proof of ownership. A duplicate delayed clear for - # lease A can otherwise erase room state that lease B published - # immediately before B atomically binds its receive identity. - if current is None or current.generation != generation: - return False - self._receive_transfer_leases.pop(room, None) - self.request_status.pop(room, None) - self.required_prefill_response_num_table.pop(room, None) - self.prefill_response_tracker.pop(room, None) - return True - def set_kv_buffer_tensors( self, k_buffers: list, @@ -513,8 +402,7 @@ def _send_chunk_ready(self, req, chunk_idx, kv_chunk, prefill_unique_rank): self._send_multipart_locked( na.to_tcp(), [ - MOONCAKE_CHUNK_READY_WIRE_TAG, - req.transfer_generation.encode("ascii"), + b"CHUNK_READY", str(req.room).encode("ascii"), str(chunk_idx).encode("ascii"), str(kv_chunk.index_slice.start).encode("ascii"), @@ -959,9 +847,9 @@ def _validate_envelope_kv_layout( def _await_transfer_futures(self, futures) -> int: """Await a chunk's per-layer RDMA writes; return the first non-zero status. - cancel() is a no-op for a running future, so every running future is - drained before returning. No write may outlive the call that publishes - the drain proof, regardless of the legacy compatibility field.""" + cancel() is a no-op for a running future, so with deferred release on we + still drain the running ones before returning (no write may outlive this + call, which the drain-ack relies on). Off: original early-return.""" ret = 0 for future in concurrent.futures.as_completed(futures): try: @@ -972,6 +860,8 @@ def _await_transfer_futures(self, futures) -> int: ret = status for f in futures: f.cancel() + if not self.enable_deferred_decode_kv_release: + return ret return ret def send_kvcache( @@ -1270,13 +1160,12 @@ def send_aux( req: TransferInfo, prefill_aux_index: int, dst_aux_ptrs: list[int], - prefill_rank: int, ): # TODO(shangming): Fix me when nvlink_transport of Mooncake is bug-free if ( self.enable_custom_mem_pool and self.custom_mem_pool_type == "NVLINK" ) or envs.SGLANG_MOONCAKE_SEND_AUX_TCP.get(): - return self.send_aux_tcp(req, prefill_aux_index, dst_aux_ptrs, prefill_rank) + return self.send_aux_tcp(req, prefill_aux_index, dst_aux_ptrs) transfer_blocks = [] prefill_aux_ptrs = self.kv_args.aux_data_ptrs @@ -1295,7 +1184,6 @@ def send_aux_tcp( req: TransferInfo, prefill_aux_index: int, dst_aux_ptrs: list[int], - prefill_rank: int, ): prefill_aux_ptrs = self.kv_args.aux_data_ptrs prefill_aux_item_lens = self.kv_args.aux_item_lens @@ -1309,8 +1197,6 @@ def send_aux_tcp( remote=req.endpoint, dst_port=req.dst_port, room=req.room, - transfer_generation=req.transfer_generation, - prefill_rank=prefill_rank, buffer_index=i, aux_index=req.dst_aux_index, data=data, @@ -1323,8 +1209,6 @@ def send_aux_data_to_endpoint( remote: str, dst_port: int, room: int, - transfer_generation: str, - prefill_rank: int, buffer_index: int, aux_index: int, data: bytes, @@ -1334,9 +1218,7 @@ def send_aux_data_to_endpoint( na.to_tcp(), [ MooncakeKVManager.AUX_DATA_HEADER, - validate_transfer_generation(transfer_generation).encode("ascii"), str(room).encode("ascii"), - str(prefill_rank).encode("ascii"), str(buffer_index).encode("ascii"), str(aux_index).encode("ascii"), struct.pack(">I", len(data)), @@ -1345,62 +1227,25 @@ def send_aux_data_to_endpoint( is_ipv6=na.is_ipv6, ) - def _handle_aux_data(self, msg: List[bytes]) -> bool: - """Handle one exact-lease AUX_DATA message without room-reuse ABA.""" - try: - if len(msg) != 8 or msg[0] != MOONCAKE_AUX_DATA_WIRE_TAG: - raise ValueError("unsupported AUX_DATA wire layout") - generation = validate_transfer_generation(msg[1].decode("ascii")) - room = _parse_mooncake_control_uint(msg[2], "room") - prefill_rank = _parse_mooncake_control_uint(msg[3], "prefill rank") - buffer_index = _parse_mooncake_control_uint(msg[4], "buffer index") - aux_index = _parse_mooncake_control_uint(msg[5], "aux index") - if type(msg[6]) is not bytes or len(msg[6]) != 4: - raise ValueError("AUX_DATA length must be one uint32 frame") - data_length = struct.unpack(">I", msg[6])[0] - data = msg[7] - if type(data) is not bytes or len(data) != data_length: - raise ValueError("AUX_DATA payload length mismatch") - except ( - AttributeError, - IndexError, - TypeError, - UnicodeError, - ValueError, - ) as error: - logger.warning("Dropping malformed or legacy AUX_DATA: %s", error) - return False - - # Identity validation and the destination write are one critical - # section with receiver clear/rebind. An old lease can therefore never - # pass validation and then write through a newly reused room. - with self._abort_state_lock: - lease = self._receive_transfer_leases.get(room) - if not self._receive_control_matches_unlocked( - lease, generation, prefill_rank - ): - logger.warning( - "Dropping stale or unexpected-rank AUX_DATA room=%s " - "generation=%s rank=%s", - room, - generation, - prefill_rank, - ) - return False - try: - AuxDataCodec.deserialize_data_to_buffer( - self.kv_args, buffer_index, aux_index, data - ) - except Exception as error: - logger.error("Failed to apply authenticated AUX_DATA: %s", error) - return False + def _handle_aux_data(self, msg: List[bytes]): + """Handle AUX_DATA messages received by the decode thread.""" + room = int(msg[1].decode("ascii")) + buffer_index = int(msg[2].decode("ascii")) + aux_index = int(msg[3].decode("ascii")) + data_length = struct.unpack(">I", msg[4])[0] + data = msg[5] + + if len(data) != data_length: + logger.error(f"AUX_DATA length mismatch for bootstrap_room {room}") + return + + AuxDataCodec.deserialize_data_to_buffer( + self.kv_args, buffer_index, aux_index, data + ) logger.debug( - "Received generation-bound AUX_DATA for bootstrap_room %s with length:%s", - room, - len(data), + f"Received AUX_DATA for bootstrap_room {room} with length:{len(data)}" ) - return True def _get_dsa_cache_transfer_skip_flags( self, info: Optional[KVArgsRegisterInfo] @@ -1839,20 +1684,12 @@ def _send_mamba_state_slice( return self._transfer_data(req.mooncake_session_id, transfer_blocks) def sync_status_to_decode_endpoint( - self, - remote: str, - dst_port: int, - room: int, - status: int, - prefill_rank: int, - transfer_generation: str, + self, remote: str, dst_port: int, room: int, status: int, prefill_rank: int ): na = NetworkAddress(remote, dst_port) self._send_multipart_locked( na.to_tcp(), [ - MOONCAKE_TRANSFER_STATUS_WIRE_TAG, - validate_transfer_generation(transfer_generation).encode("ascii"), str(room).encode("ascii"), str(status).encode("ascii"), str(prefill_rank).encode("ascii"), @@ -1890,7 +1727,7 @@ def transfer_worker( # the predicate the abort ack relies on. The flag survives # re-enqueue on defer. if not kv_chunk.staging_counted: - self._count_transfer_writer(kv_chunk.room) + self._staging_outstanding[kv_chunk.room] += 1 kv_chunk.staging_counted = True if ( @@ -1906,13 +1743,10 @@ def transfer_worker( MooncakeRequestStage.MOONCAKE_WORKER_SEND.level, thread_finish_flag=True, ) - remaining_writers = self._uncount_transfer_writer(kv_chunk.room) - kv_chunk.staging_counted = False - if remaining_writers == 0 and self.mark_transfer_quiescent( - kv_chunk.room - ): - self.transfer_infos.pop(kv_chunk.room, None) - self.req_to_decode_prefix_len.pop(kv_chunk.room, None) + self._staging_outstanding.pop(kv_chunk.room, None) + if self.enable_deferred_decode_kv_release: + # Skipped => nothing written for this aborted room; ack. + self._maybe_ack_drained_abort(kv_chunk.room) continue if ( @@ -1954,7 +1788,6 @@ def transfer_worker( req.room, KVPoll.Failed, prefill_unique_rank, - req.transfer_generation, ) break @@ -2109,7 +1942,6 @@ def transfer_worker( req.room, KVPoll.Failed, prefill_unique_rank, - req.transfer_generation, ) break @@ -2141,7 +1973,6 @@ def transfer_worker( req.room, KVPoll.Failed, prefill_unique_rank, - req.transfer_generation, ) break @@ -2150,35 +1981,23 @@ def transfer_worker( req, kv_chunk.prefill_aux_index, target_rank_registration_info.dst_aux_ptrs, - prefill_unique_rank, ) polls.append(True if ret == 0 else False) dst_ranks_infos.append( - ( - req.endpoint, - req.dst_port, - req.room, - req.transfer_generation, - ) + (req.endpoint, req.dst_port, req.room) ) # Only sync status when all the dst ranks have received the kvcache if len(polls) == req.required_dst_info_num: status = KVPoll.Success if all(polls) else KVPoll.Failed self.update_status(req.room, status) - for ( - endpoint, - dst_port, - room, - generation, - ) in dst_ranks_infos: + for endpoint, dst_port, room in dst_ranks_infos: self.sync_status_to_decode_endpoint( endpoint, dst_port, room, status, prefill_unique_rank, - generation, ) else: # Dummy request means the decode instance is not used, so its status can be marked as success directly @@ -2203,22 +2022,25 @@ def transfer_worker( if staging_deferred: continue - remaining_writers = self._uncount_transfer_writer(kv_chunk.room) - kv_chunk.staging_counted = False + self._staging_outstanding[kv_chunk.room] -= 1 + if self.enable_deferred_decode_kv_release: + # In-flight write finished; if aborted and nothing outstanding, + # the pages are idle -> release the held ack. + self._maybe_ack_drained_abort(kv_chunk.room) # Tear down only when no chunk is still outstanding and the room # has concluded: already cleared, Success, or a Failed *last* # chunk. A non-last Failed chunk keeps the room (more chunks may # follow), not on the last chunk alone since an earlier deferred # chunk may still need to transfer. - if ( - remaining_writers == 0 - and ( - kv_chunk.room not in self.request_status - or self.check_status(kv_chunk.room) == KVPoll.Success - or self.check_status(kv_chunk.room) == KVPoll.Failed + if self._staging_outstanding.get(kv_chunk.room, 0) <= 0 and ( + kv_chunk.room not in self.request_status + or self.check_status(kv_chunk.room) == KVPoll.Success + or ( + kv_chunk.is_last_chunk + and self.check_status(kv_chunk.room) == KVPoll.Failed ) - and self.mark_transfer_quiescent(kv_chunk.room) ): + self._staging_outstanding.pop(kv_chunk.room, None) if kv_chunk.room in self.transfer_infos: self.transfer_infos.pop(kv_chunk.room) self.req_to_decode_prefix_len.pop(kv_chunk.room, None) @@ -2248,52 +2070,76 @@ def bootstrap_thread(): handle_watermark_msg(self._staging_ctx, waiting_req_bytes) continue # Staging: decode replies with allocated staging offset - if waiting_req_bytes[0] == STAGING_RSP_WIRE_TAG: + 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": - try: - if len(waiting_req_bytes) != 5: - raise ValueError("ABORT must contain exactly five frames") - 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")) - transfer_generation = validate_transfer_generation( - waiting_req_bytes[4].decode("ascii") - ) - except Exception as error: - logger.warning("Ignoring malformed ABORT: %s", error) - continue - generation_active = ( - self.active_transfer_generation(room_to_be_aborted) - == transfer_generation - ) + 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 = ( - generation_active - and room_to_be_aborted in self.request_status + 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 nothing is in + # flight, decode falls back to the release timeout. + if room_active: + 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" + ) + elif self._staging_outstanding.get(room_to_be_aborted, 0) == 0: + # Concluded/unknown AND quiescent: ack now. A cleared + # room is not automatically quiescent -- clear() can + # drop a room whose chunk is still transferring. + 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" ) - if self.register_deferred_ack_target( - room_to_be_aborted, - transfer_generation, - decode_ip, - decode_port, - ): - self._maybe_ack_drained_abort( - room_to_be_aborted, transfer_generation + 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, ) - elif not generation_active: logger.debug( - "Ignoring stale ABORT for room %s generation %s", - room_to_be_aborted, - transfer_generation, + 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}" ) continue mooncake_session_id = waiting_req_bytes[3].decode("ascii") @@ -2322,37 +2168,14 @@ def bootstrap_thread(): ) continue else: - try: - required_dst_info_num = int( - waiting_req_bytes[7].decode("ascii") - ) - room = int(room) - transfer_info = TransferInfo.from_zmq(waiting_req_bytes) - except (IndexError, UnicodeError, ValueError) as error: - # A recyclable room number is not an identity. Missing - # or malformed generation metadata cannot fail whichever - # newer lease happens to own that number, and must not - # terminate the long-lived listener thread. - logger.warning( - "Dropping unauthenticated Mooncake room metadata: %s", - error, - ) - continue - if not self.register_transfer_generation( - room, transfer_info.transfer_generation - ): - logger.warning( - "Dropping late Mooncake metadata for room %s " - "generation %s; active generation is %s", - room, - transfer_info.transfer_generation, - self.active_transfer_generation(room), - ) - continue + required_dst_info_num = int(waiting_req_bytes[7].decode("ascii")) + room = int(room) if room not in self.transfer_infos: self.transfer_infos[room] = {} - self.transfer_infos[room][mooncake_session_id] = transfer_info + self.transfer_infos[room][mooncake_session_id] = ( + TransferInfo.from_zmq(waiting_req_bytes) + ) # NOTE: after bootstrapping we can mark the req as waiting for input if len(self.transfer_infos[room]) == required_dst_info_num: self.resolve_kv_replica_factor(self.transfer_infos[room]) @@ -2377,152 +2200,72 @@ def decode_thread(): continue # Staging: prefill notifies a chunk written to staging buffer - if msg[0] == MOONCAKE_CHUNK_READY_WIRE_TAG: - self._handle_chunk_ready(msg) + if msg[0] == b"CHUNK_READY": + room = int(msg[1].decode("ascii")) + chunk_idx = int(msg[2].decode("ascii")) + page_start = int(msg[3].decode("ascii")) + num_pages = int(msg[4].decode("ascii")) + session_id = msg[5].decode("ascii") + handler = self._staging_handler + assert handler is not None, ( + "CHUNK_READY received before staging handler initialized" + ) + handler.handle_chunk_arrived( + room, + chunk_idx, + page_start, + num_pages, + session_id, + ) continue # Staging: prefill pre-requests staging allocation before forward - if msg[0] == STAGING_REQ_WIRE_TAG: + if msg[0] == b"STAGING_REQ": self._handle_staging_req(msg) continue # Prefill acknowledges abort notification if msg[0] == b"ABORT_ACK": - try: - if len(msg) != 4: - raise ValueError( - "ABORT_ACK must contain exactly four frames" - ) - ack_aborted_room = int(msg[1].decode("ascii")) - prefill_rank = int(msg[2].decode("ascii")) - generation = validate_transfer_generation( - msg[3].decode("ascii") - ) - except (UnicodeDecodeError, ValueError): - logger.warning( - "Ignoring generation-less or malformed 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")) ) - else: - logger.debug("Received ABORT_ACK for room %s", ack_aborted_room) - self.note_abort_ack(ack_aborted_room, generation, prefill_rank) continue - if msg[0] == MOONCAKE_TRANSFER_STATUS_WIRE_TAG: - self._handle_transfer_status(msg) - continue - logger.warning( - "Dropping unknown, legacy, or malformed Mooncake control frame: %r", - msg[0][:32] if msg else b"", - ) + bootstrap_room, status, prefill_rank = msg + status = int(status.decode("ascii")) + bootstrap_room = int(bootstrap_room.decode("ascii")) + prefill_rank = int(prefill_rank.decode("ascii")) + + if status == KVPoll.Success: + if bootstrap_room in self.request_status: + self.prefill_response_tracker[bootstrap_room].add(prefill_rank) + expected_response_num = ( + self.required_prefill_response_num_table[bootstrap_room] + ) + arrived_response_num = len( + self.prefill_response_tracker[bootstrap_room] + ) + if arrived_response_num == expected_response_num: + if self.enable_staging: + handler = self._staging_handler + if handler.is_staging_room(bootstrap_room): + handler.submit_last_scatter_async(bootstrap_room) + self.update_status(bootstrap_room, KVPoll.Success) + elif status == KVPoll.Failed: + self.record_failure( + bootstrap_room, + "Failed to get kvcache from prefill instance, it might be dead", + ) + self.update_status(bootstrap_room, status) threading.Thread(target=decode_thread).start() self._start_heartbeat_checker_thread() - def _handle_chunk_ready(self, msg: List[bytes]) -> bool: - try: - if len(msg) != 8 or msg[0] != MOONCAKE_CHUNK_READY_WIRE_TAG: - raise ValueError("unsupported CHUNK_READY wire layout") - generation = validate_transfer_generation(msg[1].decode("ascii")) - room = _parse_mooncake_control_uint(msg[2], "room") - chunk_idx = _parse_mooncake_control_uint(msg[3], "chunk index") - page_start = _parse_mooncake_control_uint(msg[4], "page start") - num_pages = _parse_mooncake_control_uint( - msg[5], "page count", positive=True - ) - session_id = msg[6].decode("ascii") - if not session_id: - raise ValueError("CHUNK_READY session must be non-empty") - prefill_rank = _parse_mooncake_control_uint(msg[7], "prefill rank") - except ( - AttributeError, - IndexError, - TypeError, - UnicodeError, - ValueError, - ) as error: - logger.warning("Dropping malformed or legacy CHUNK_READY: %s", error) - return False - - with self._abort_state_lock: - lease = self._receive_transfer_leases.get(room) - if not self._receive_control_matches_unlocked( - lease, - generation, - prefill_rank, - session_id=session_id, - ): - logger.warning( - "Dropping stale or unexpected-rank CHUNK_READY room=%s " - "generation=%s rank=%s", - room, - generation, - prefill_rank, - ) - return False - handler = self._staging_handler - if handler is None: - logger.warning("CHUNK_READY received before staging initialization") - return False - handler.handle_chunk_arrived( - room, - chunk_idx, - page_start, - num_pages, - prefill_rank, - ) - return True - - def _handle_transfer_status(self, msg: List[bytes]) -> bool: - try: - if len(msg) != 5 or msg[0] != MOONCAKE_TRANSFER_STATUS_WIRE_TAG: - raise ValueError("unsupported TRANSFER_STATUS wire layout") - generation = validate_transfer_generation(msg[1].decode("ascii")) - room = _parse_mooncake_control_uint(msg[2], "room") - status = _parse_mooncake_control_uint(msg[3], "status") - prefill_rank = _parse_mooncake_control_uint(msg[4], "prefill rank") - if status not in (KVPoll.Success, KVPoll.Failed): - raise ValueError("TRANSFER_STATUS must be Success or Failed") - except ( - AttributeError, - IndexError, - TypeError, - UnicodeError, - ValueError, - ) as error: - logger.warning("Dropping malformed or legacy TRANSFER_STATUS: %s", error) - return False - - with self._abort_state_lock: - lease = self._receive_transfer_leases.get(room) - if not self._receive_control_matches_unlocked( - lease, generation, prefill_rank - ): - logger.warning( - "Dropping stale or unexpected-rank TRANSFER_STATUS room=%s " - "generation=%s rank=%s", - room, - generation, - prefill_rank, - ) - return False - if status == KVPoll.Failed: - self.record_failure( - room, - "Failed to get kvcache from prefill instance, it might be dead", - ) - self.update_status(room, KVPoll.Failed) - return True - - self.prefill_response_tracker[room].add(prefill_rank) - if self.prefill_response_tracker[room] != lease.expected_prefill_ranks: - return True - if self.enable_staging: - handler = self._staging_handler - if handler is not None and handler.is_staging_room(room): - handler.submit_last_scatter_async(room) - self.update_status(room, KVPoll.Success) - return True - def add_transfer_request( self, bootstrap_room: int, @@ -2756,46 +2499,10 @@ def __init__( mgr: MooncakeKVManager, bootstrap_addr: str, bootstrap_room: Optional[int] = None, - transfer_generation: Optional[str] = None, ): self.session_id = mgr.get_session_id() self.init_time = None - super().__init__( - mgr, - bootstrap_addr, - bootstrap_room, - transfer_generation=transfer_generation, - ) - # Publish the generation before any delayed control can be attributed - # to the freshly created room. Rank identity is bound after bootstrap. - self.kv_mgr._register_receive_transfer_lease( - self.bootstrap_room, - self.transfer_generation, - self.session_id, - ) - - def init(self, prefill_dp_rank: int): - super().init(prefill_dp_rank) - if self.conclude_state == KVPoll.Failed: - return - expected_prefill_ranks = frozenset( - bootstrap_info["prefill_unique_rank"] - for bootstrap_info in self.bootstrap_infos - if not bootstrap_info["is_dummy"] - ) - self.kv_mgr._register_receive_transfer_lease( - self.bootstrap_room, - self.transfer_generation, - self.session_id, - expected_prefill_ranks, - ) - - def clear(self) -> None: - # The exact-generation compare and all room-keyed state removal share - # one lock; a delayed clear from lease A cannot erase lease B. - self.kv_mgr._clear_receive_transfer_lease( - self.bootstrap_room, self.transfer_generation - ) + super().__init__(mgr, bootstrap_addr, bootstrap_room) def _register_kv_args(self) -> bool: for bootstrap_info in self.bootstrap_infos: @@ -2937,7 +2644,6 @@ def send_metadata( if not is_dummy and device_kv_indices is not None else b"" ), - self.transfer_generation.encode("ascii"), ] ) except zmq.ZMQError: diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 463d1eb9badf..ff5f3d7c9b9d 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -8,7 +8,7 @@ import time import uuid from collections import defaultdict -from typing import TYPE_CHECKING, Any, Dict, FrozenSet, List, Optional, Set, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple import numpy as np import numpy.typing as npt @@ -24,12 +24,9 @@ CommonKVReceiver, CommonKVSender, KVTransferError, - validate_transfer_generation, ) from sglang.srt.disaggregation.common.staging_handler import ( STAGING_WATERMARK_WAIT_S, - STAGING_REQ_WIRE_TAG, - STAGING_RSP_WIRE_TAG, StagingManagerMixin, handle_staging_rsp, handle_watermark_msg, @@ -74,126 +71,6 @@ GUARD = "NixlMsgGuard".encode("ascii") KV_MEM_KINDS = {"VRAM", "DRAM"} -_COMPLETION_NOTIFICATION_VERSION = "nixlv2" -_MAX_NOTIFICATION_UINT_DIGITS = 20 - - -def _parse_notification_uint(value: str, field: str, *, positive: bool = False) -> int: - """Parse one canonical bounded unsigned integer from a completion tag.""" - - if ( - not value - or len(value) > _MAX_NOTIFICATION_UINT_DIGITS - or not value.isascii() - or not value.isdecimal() - or (len(value) > 1 and value[0] == "0") - ): - raise ValueError(f"{field} must be canonical unsigned decimal") - parsed = int(value) - if positive and parsed == 0: - raise ValueError(f"{field} must be positive") - return parsed - - -def _parse_completion_notification( - message: bytes, -) -> Tuple[str, int, str, Tuple[Any, ...]]: - """Parse the strict, generation-bound NIXL completion wire contract. - - Every accepted tag starts with ``nixlv2_{generation}_{room}_{kind}``. - Legacy/generation-less tags are intentionally incompatible: accepting one - would let a queued completion from a retired room lease complete its reuse. - """ - - if type(message) is not bytes: - raise ValueError("completion notification must be exact bytes") - components = message.decode("ascii").split("_") - if len(components) < 4 or components[0] != _COMPLETION_NOTIFICATION_VERSION: - raise ValueError("unsupported completion notification version") - generation = validate_transfer_generation(components[1]) - room = _parse_notification_uint(components[2], "room") - kind = components[3] - - if kind == "kv": - if len(components) not in (7, 10): - raise ValueError("invalid kv completion field count") - chunk_id = _parse_notification_uint(components[4], "chunk_id") - if components[5] not in ("0", "1"): - raise ValueError("is_last must be exactly 0 or 1") - is_last = components[5] == "1" - pp_rank = _parse_notification_uint(components[6], "pp_rank") - if len(components) == 7: - return generation, room, kind, (chunk_id, is_last, pp_rank) - if components[7] != "part": - raise ValueError("invalid kv-part discriminator") - part_idx = _parse_notification_uint(components[8], "part_idx") - num_parts = _parse_notification_uint(components[9], "num_parts", positive=True) - if part_idx >= num_parts: - raise ValueError("kv-part index is outside its declared part count") - return ( - generation, - room, - "kvpart", - ( - chunk_id, - is_last, - pp_rank, - part_idx, - num_parts, - ), - ) - - if kind == "state": - if len(components) not in (5, 6): - raise ValueError("invalid state completion field count") - pp_rank = _parse_notification_uint(components[4], "pp_rank") - component_id = ( - _parse_notification_uint(components[5], "state_component") - if len(components) == 6 - else None - ) - return generation, room, kind, (pp_rank, component_id) - - if kind == "aux": - if len(components) == 5: - pp_rank = _parse_notification_uint(components[4], "pp_rank") - return generation, room, kind, (pp_rank, None) - if len(components) != 7 or components[4] != "nokv": - raise ValueError("invalid aux completion layout") - pp_rank = _parse_notification_uint(components[5], "pp_rank") - expected = _parse_notification_uint(components[6], "expected_kv_chunks") - return generation, room, kind, (pp_rank, expected) - - if kind == "stg": - if len(components) < 11: - raise ValueError("invalid staging completion field count") - chunk_id = _parse_notification_uint(components[4], "chunk_id") - if components[5] not in ("0", "1"): - raise ValueError("is_last must be exactly 0 or 1") - is_last = components[5] == "1" - pp_rank = _parse_notification_uint(components[6], "pp_rank") - chunk_idx = _parse_notification_uint(components[7], "staging_chunk_idx") - page_start = _parse_notification_uint(components[8], "page_start") - num_pages = _parse_notification_uint(components[9], "num_pages", positive=True) - agent_name = "_".join(components[10:]) - if not agent_name: - raise ValueError("staging agent name must be non-empty") - return ( - generation, - room, - kind, - ( - chunk_id, - is_last, - pp_rank, - chunk_idx, - page_start, - num_pages, - agent_name, - ), - ) - - raise ValueError(f"unsupported completion notification kind: {kind!r}") def _normalize_kv_mem_kinds(kinds: Optional[List[str]], expected_len: int) -> List[str]: @@ -290,23 +167,10 @@ class TransferInfo: dst_state_indices: List[List[int]] decode_prefix_len: Optional[int] = None # for decode radix cache is_dummy_rank: Optional[bool] = None - transfer_generation: str = "" # NOTE: optional staging field; populated via STAGING_RSP. Keep at the # end so positional construction in from_zmq() continues to work. staging: Optional[StagingTransferInfo] = None - def __post_init__(self): - generation = validate_transfer_generation(self.transfer_generation) - # Cache the immutable per-lease prefix once on metadata ingestion. The - # per-chunk hot path only appends its small kind-specific suffix. - self._completion_notification_prefix = ( - f"{_COMPLETION_NOTIFICATION_VERSION}_{generation}_{self.room}" - ) - - @property - def completion_notification_prefix(self) -> str: - return self._completion_notification_prefix - def is_dummy(self): # A transfer is "dummy" only for CP non-authoritative ranks. # When dst_kv_indices is empty due to a decode-side radix cache @@ -341,9 +205,6 @@ def from_zmq(cls, msg: List[bytes]): if len(msg) > 9 and msg[9] != b"" else None ), - transfer_generation=validate_transfer_generation( - msg[10].decode("ascii") if len(msg) > 10 else None - ), ) @@ -499,13 +360,6 @@ def repeat_indices_over_layers( class TransferStatus: """Used by KV Receiver to know when a transfer is done.""" - # Exact receiver lease identity. A room number alone is recyclable and - # therefore cannot authenticate an attached completion notification. - transfer_generation: str - # Bound after bootstrap discovery and before metadata submission. Every - # completion family must name one exact non-dummy remote writer. - expected_source_ranks: Optional[FrozenSet[int]] = None - # KV chunks received per pp_rank: {pp_rank: set of chunk_ids} received_kvs_per_pp: Dict[int, Set[int]] = dataclasses.field( default_factory=lambda: defaultdict(set) @@ -525,51 +379,16 @@ class TransferStatus: received_kv_parts_per_pp: Optional[Dict[Tuple[int, int], Set[int]]] = None expected_kv_parts_per_pp: Optional[Dict[Tuple[int, int], int]] = None - def __post_init__(self): - self.transfer_generation = validate_transfer_generation( - self.transfer_generation - ) - if self.expected_source_ranks is not None: - self.bind_expected_source_ranks(self.expected_source_ranks) - - def bind_expected_source_ranks(self, ranks) -> None: - expected = frozenset(ranks) - if not expected or any(type(rank) is not int or rank < 0 for rank in expected): - raise ValueError( - "expected completion ranks must be non-empty non-negative ints" - ) - if ( - self.expected_source_ranks is not None - and self.expected_source_ranks != expected - ): - raise RuntimeError("completion source-rank identity set is immutable") - self.expected_source_ranks = expected - self.num_pp_ranks_expected = len(expected) - - def accepts_source_rank(self, rank: object) -> bool: - return ( - type(rank) is int - and self.expected_source_ranks is not None - and rank in self.expected_source_ranks - ) - def is_done(self): if self.num_pp_ranks_expected is None or not self.received_aux: return False # If state data is expected, check all PP ranks have sent it - if self.expects_state: - if self.expected_source_ranks is not None: - if not self.expected_source_ranks.issubset(self.received_state_per_pp): - return False - elif len(self.received_state_per_pp) < self.num_pp_ranks_expected: - # Compatibility for direct status construction outside the - # receiver path. Production status is exact-set bound in init(). - return False - # All PP ranks must have reported their expected count - if self.expected_source_ranks is not None and set( - self.expected_kvs_per_pp - ) != set(self.expected_source_ranks): + if ( + self.expects_state + and len(self.received_state_per_pp) < self.num_pp_ranks_expected + ): return False + # All PP ranks must have reported their expected count if len(self.expected_kvs_per_pp) < self.num_pp_ranks_expected: return False # Each PP rank must have received all expected chunks @@ -588,9 +407,9 @@ def __init__( is_mla_backend: Optional[bool] = False, ): super().__init__(args, disaggregation_mode, server_args, is_mla_backend) - # One collision-free TP×PP×CP identity, shared with ABORT_ACK and the - # receiver's authenticated bootstrap-derived source-rank set. - self.transfer_source_rank = self._prefill_unique_rank() + self.transfer_source_rank = ( + self.kv_args.pp_rank * get_parallel().tp_size + self.kv_args.engine_rank + ) self.kv_args.kv_data_mem_kinds = _normalize_kv_mem_kinds( getattr(self.kv_args, "kv_data_mem_kinds", None), len(self.kv_args.kv_data_ptrs), @@ -601,10 +420,15 @@ def __init__( and self.kv_args.kv_data_mem_kinds else None ) - self._use_torch_transfer = ( - envs.SGLANG_DISAGGREGATION_NIXL_USE_TORCH_TRANSFER.get() - ) - self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() + try: + from nixl._api import nixl_agent, nixl_agent_config, nixl_thread_sync_t + except ImportError as e: + raise ImportError( + "Please install NIXL by following the instructions at " + "https://github.com/ai-dynamo/nixl/blob/main/README.md " + "to run SGLang with NixlTransferEngine." + ) from e + backend = envs.SGLANG_DISAGGREGATION_NIXL_BACKEND.get() num_threads = 8 if disaggregation_mode == DisaggregationMode.PREFILL else 0 backend_params = json.loads( @@ -618,6 +442,19 @@ def __init__( "SGLANG_DISAGGREGATION_NIXL_BACKEND_PARAMS must be a JSON object " "with string keys and string values" ) + # self.transfer_worker and self._start_bootstrap_thread runs concurrently + # so we cannot use sync_mode=None which is thread-unsafe. + agent_config = nixl_agent_config( + backends=[], + num_threads=num_threads, + sync_mode=nixl_thread_sync_t.NIXL_THREAD_SYNC_STRICT, + ) + self.agent = nixl_agent(str(uuid.uuid4()), agent_config) + if envs.SGLANG_DISAGGREGATION_NIXL_USE_TORCH_TRANSFER.get(): + from nixl.torch_transfer import TorchTransferAgent + + # The manager's allocation lifetime contract remains unchanged. + self.agent = TorchTransferAgent(self.agent, owner=self) if num_threads > 0: # TODO: Remove this once NIXL passes thread parameters from # nixl_agent_config to explicitly-created backends. @@ -627,55 +464,6 @@ def __init__( backend_params.setdefault("thread_count", str(num_threads)) elif backend == "UCCL": backend_params.setdefault("num_cpus", str(num_threads)) - - agent_name = str(uuid.uuid4()) - if self._use_torch_transfer: - if backend != "UCX": - raise NotImplementedError( - "Experimental PyTorch transfer currently supports only " - "the NIXL UCX backend" - ) - if self.enable_staging: - raise NotImplementedError( - "Experimental PyTorch transfer does not yet support " - "SGLang staging buffers" - ) - if len(set(self.kv_args.kv_data_mem_kinds)) > 1: - raise NotImplementedError( - "Experimental PyTorch transfer does not yet support mixed " - "KV memory kinds" - ) - from sglang.srt.disaggregation.nixl.torch_transfer import ( - TorchTransferAgent, - ) - - self.agent = TorchTransferAgent( - agent_name, - owner=self, - background_progress=True, - num_threads=num_threads, - ) - else: - try: - from nixl._api import ( - nixl_agent, - nixl_agent_config, - nixl_thread_sync_t, - ) - except ImportError as e: - raise ImportError( - "Please install NIXL by following the instructions at " - "https://github.com/ai-dynamo/nixl/blob/main/README.md " - "to run SGLang with NixlTransferEngine." - ) from e - - # Bootstrap and transfer-worker threads share the agent. - agent_config = nixl_agent_config( - backends=[], - num_threads=num_threads, - sync_mode=nixl_thread_sync_t.NIXL_THREAD_SYNC_STRICT, - ) - self.agent = nixl_agent(agent_name, agent_config) run_with_deadline( lambda: self.agent.create_backend(backend, backend_params), timeout_s=envs.SGLANG_DISAGGREGATION_ENGINE_INIT_TIMEOUT.get(), @@ -688,11 +476,11 @@ def __init__( f"NIXL backend '{backend}' not found. Available: {available_plugins}. " f"Please install the required NIXL plugin or choose from: {available_plugins}" ) - data_plane = "PyTorch transfer" if self._use_torch_transfer else "native NIXL" - logger.info(f"NIXL KVManager initialized with {data_plane} backend: {backend}") + logger.info(f"NIXL KVManager initialized with backend: {backend}") self.register_buffer_to_engine() + self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() self.kv_buffer_tensors = None self.prep_handles: Dict[str, Any] = {} self.prep_handle_slice_src: Optional[Tuple[Any, int, int, int]] = ( @@ -713,10 +501,6 @@ def __init__( FastQueue() for _ in range(transfer_queue_size) ] self.exceptions: Dict[int, Exception] = {} - # Peer registration happens before a room-specific bootstrap - # message. Preserve failures by remote agent so the later room can - # transition to Failed instead of waiting after the listener dies. - self._bootstrap_errors: Dict[str, 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) @@ -739,14 +523,14 @@ def __init__( ).start() self._start_bootstrap_thread() elif self.disaggregation_mode == DisaggregationMode.DECODE: - # Receiver construction binds each entry to an exact generation. - # Never auto-create room-only status from an untrusted notification. - self.transfer_statuses: Dict[int, TransferStatus] = {} + self.transfer_statuses: Dict[int, TransferStatus] = defaultdict( + TransferStatus + ) if self.enable_staging: self._init_staging_decode_ctx() self._staging_handler = None - # ABORT_ACK is mandatory even without heterogeneous-TP staging. - self._start_decode_listener_thread() + if self.enable_staging or self.enable_deferred_decode_kv_release: + self._start_decode_listener_thread() self._start_heartbeat_checker_thread() else: raise ValueError( @@ -823,27 +607,15 @@ def _start_decode_listener_thread(self): def decode_listener_thread(): while True: msg = self.server_socket.recv_multipart() - if msg[0] == STAGING_REQ_WIRE_TAG: + if msg[0] == b"STAGING_REQ": self._handle_staging_req(msg) continue if msg[0] == b"ABORT_ACK": - # Only a complete exact-generation ACK can release memory. - try: - if len(msg) != 4: - raise ValueError( - "ABORT_ACK must contain exactly four frames" - ) - room = int(msg[1].decode("ascii")) - prefill_rank = int(msg[2].decode("ascii")) - generation = validate_transfer_generation( - msg[3].decode("ascii") - ) - except (UnicodeDecodeError, ValueError): - logger.warning( - "Ignoring generation-less or malformed 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")) ) - else: - self.note_abort_ack(room, generation, prefill_rank) continue logger.warning( "decode_listener_thread: unexpected message tag %s", @@ -936,27 +708,16 @@ def _prep_equal_tp_dlist( f"xfer_len={xfer_len}, item_len={item_len}, mem_kind={mem_kind}" ) n = num_slots if num_slots is not None else (data_len // item_len) - if self._use_torch_transfer: - # Core indexed plans flatten the logical blocks represented by - # each strided Region. Keep one compact row per allocation; - # existing layer-major/slot-minor indices remain unchanged. - arrays.append( - np.asarray( - [[base_ptr, xfer_len, device_id, item_len, n]], - dtype=np.uint64, - ) - ) - else: - addrs = np.arange(n, dtype=np.uint64) * np.uint64(item_len) + base_ptr - arrays.append( - np.column_stack( - [ - addrs, - np.full(n, xfer_len, dtype=np.uint64), - np.full(n, device_id, dtype=np.uint64), - ] - ) + addrs = np.arange(n, dtype=np.uint64) * np.uint64(item_len) + base_ptr + arrays.append( + np.column_stack( + [ + addrs, + np.full(n, xfer_len, dtype=np.uint64), + np.full(n, device_id, dtype=np.uint64), + ] ) + ) prep_handle = self.agent.prep_xfer_dlist(peer_name, np.vstack(arrays), mem_kind) assert prep_handle is not None, ( @@ -1339,17 +1100,13 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) kv_chunk: TransferKVChunk = queue.get() room = kv_chunk.room handles: List[Any] = [] - batch_active = False try: - if self._use_torch_transfer: - self.agent.begin_handle_batch(handles) - batch_active = True if room not in self.request_status: logger.debug( "Skipping chunk for room %s because it has been cleared", room, ) - self._maybe_ack_drained_abort(room) + self._staging_outstanding.pop(room, None) continue # Counted at dequeue, before the status check, so @@ -1357,15 +1114,14 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) # the predicate the abort ack relies on. The flag survives # re-enqueue on defer. if not kv_chunk.staging_counted: - self._count_transfer_writer(room) + self._staging_outstanding[room] += 1 kv_chunk.staging_counted = True if self.check_status(room) == KVPoll.Failed: - remaining_writers = self._uncount_transfer_writer(room) - kv_chunk.staging_counted = False - if remaining_writers == 0 and self.mark_transfer_quiescent(room): - self.transfer_infos.pop(room, None) - self.req_to_decode_prefix_len.pop(room, None) + self._staging_outstanding.pop(room, None) + if self.enable_deferred_decode_kv_release: + # Skipped => nothing written for this aborted room; ack. + self._maybe_ack_drained_abort(room) continue room_transfer_infos = self.transfer_infos.get(room) @@ -1375,25 +1131,9 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) "has been cleared", room, ) - remaining_writers = self._uncount_transfer_writer(room) - kv_chunk.staging_counted = False - if remaining_writers == 0 and self.check_status(room) in ( - KVPoll.Success, - KVPoll.Failed, - ): - self.mark_transfer_quiescent(room) + self._staging_outstanding.pop(room, None) continue - # Early-send can be queued while the preceding forward is still - # writing the source KV pages. Consume the event once, before - # any path can read those pages (staging, DCP packing, or direct - # NIXL/PyTorch transfer). Clear it before waiting so a staging - # defer/re-enqueue cannot repeat the host synchronization. - wait_event = kv_chunk.wait_event - kv_chunk.wait_event = None - if wait_event is not None: - wait_event.synchronize() - # Lazily build a per-worker staging strategy bound to this # worker's private staging buffer (matches mooncake). if ( @@ -1458,7 +1198,7 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) src_prefill_kv_indices = kv_chunk.prefill_kv_indices notif = ( - f"{req.completion_notification_prefix}_kv_{kv_chunk.chunk_id}" + f"{req.room}_kv_{kv_chunk.chunk_id}" f"_{int(kv_chunk.is_last_chunk)}_{self.transfer_source_rank}" ) @@ -1553,12 +1293,13 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) ), ) else: - self.send_kvcache_mixed( - req.agent_name, - src_prefill_kv_indices, - chunked_dst_kv_indice, - notif, - handles, + handles.extend( + self.send_kvcache_mixed( + req.agent_name, + src_prefill_kv_indices, + chunked_dst_kv_indice, + notif, + ) ) else: kv_xfer_handle = self.send_kvcache_slice( @@ -1568,26 +1309,27 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) notif, ) - if kv_xfer_handle is not None and not self._use_torch_transfer: + if kv_xfer_handle is not None: handles.append(kv_xfer_handle) if kv_chunk.is_last_chunk: dst_info = self.decode_kv_args_table[req.agent_name] if kv_chunk.state_indices: - self.maybe_send_extra( + state_xfer_handles = self.maybe_send_extra( req.agent_name, kv_chunk.state_indices, dst_info.dst_state_data_ptrs, req.dst_state_indices, dst_info.gpu_id, - f"{req.completion_notification_prefix}_state_" - f"{self.transfer_source_rank}", + f"{req.room}_state_{self.transfer_source_rank}", decode_tp_size, decode_tp_rank=dst_info.decode_tp_rank, dst_state_item_lens=dst_info.dst_state_item_lens, dst_state_dim_per_tensor=dst_info.dst_state_dim_per_tensor, dst_state_layer_ids=dst_info.dst_state_layer_ids, - handles=handles, + ) + handles.extend( + h for h in state_xfer_handles if h is not None ) if kv_chunk.prefill_aux_index is None: @@ -1595,18 +1337,13 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) # A no-KV notification still identifies its PP source. # Empty non-final chunks do not consume chunk IDs, so a # final no-KV chunk_id equals the prior KV chunk count. + aux_notif = f"{req.room}_aux" if ( len(kv_chunk.prefill_kv_indices) == 0 or not self.kv_args.kv_data_ptrs ): - aux_notif = ( - f"{req.completion_notification_prefix}_aux_nokv_" - f"{self.transfer_source_rank}_{kv_chunk.chunk_id}" - ) - else: - aux_notif = ( - f"{req.completion_notification_prefix}_aux_" - f"{self.transfer_source_rank}" + aux_notif += ( + f"_nokv_{self.transfer_source_rank}_{kv_chunk.chunk_id}" ) aux_xfer_handle = self.send_aux( req.agent_name, @@ -1615,29 +1352,17 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) req.dst_aux_index, aux_notif, ) - if not self._use_torch_transfer: - handles.append(aux_xfer_handle) + handles.append(aux_xfer_handle) if staging_deferred: # Chunk has been re-enqueued; do not advance status. continue while handles: - if self._use_torch_transfer: - # Advance provider progress once. The following loop is - # the sole status scan for caller-owned batch Works. - self.agent.progress() all_done = True for handle in handles: state = self.agent.check_xfer_state(handle) if state == "ERR": - deferred_error = ( - self.agent.pop_xfer_error(handle) - if self._use_torch_transfer - else None - ) - if deferred_error is not None: - raise deferred_error raise RuntimeError( f"NIXL transfer encountered ERR room={room}" ) @@ -1647,8 +1372,11 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) break time.sleep(0) - remaining_writers = self._uncount_transfer_writer(room) - kv_chunk.staging_counted = False + self._staging_outstanding[room] -= 1 + if self.enable_deferred_decode_kv_release: + # Handles all DONE => this room's writes landed; ack if it + # was aborted and nothing else is outstanding. + self._maybe_ack_drained_abort(room) if kv_chunk.is_last_chunk: self.update_status(room, KVPoll.Success) elif self.check_status(room) != KVPoll.Success: @@ -1656,17 +1384,18 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) # already concluded Success; don't regress the status. self.update_status(room, KVPoll.Transferring) - # Once a room is Failed, enqueue rejects all later chunks. When - # its exact outstanding count reaches zero, retain only the - # process-lifetime generation tombstone, not mutable room state. - if ( - remaining_writers == 0 - and ( - self.check_status(room) == KVPoll.Success - or self.check_status(room) == KVPoll.Failed + # Drop per-room state only when no chunk is still outstanding and + # the room has concluded: Success, or a Failed *last* chunk. A + # non-last Failed chunk keeps the room (more chunks may follow); a + # late chunk for an already-Failed room is skipped at loop top. + if self._staging_outstanding.get(room, 0) <= 0 and ( + self.check_status(room) == KVPoll.Success + or ( + kv_chunk.is_last_chunk + and self.check_status(room) == KVPoll.Failed ) - and self.mark_transfer_quiescent(room) ): + self._staging_outstanding.pop(room, None) self.transfer_infos.pop(room, None) self.req_to_decode_prefix_len.pop(room, None) if self.enable_staging and self._staging_ctx is not None: @@ -1675,61 +1404,7 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) for k in list(self._staging_ctx.prefetch_requested): if k[0] == room: self._staging_ctx.prefetch_requested.discard(k) - except BaseException as e: - if self._use_torch_transfer and handles: - if not isinstance(e, Exception): - # One identity publication retains every handle already - # appended to the prepublished batch, even if this worker - # is about to rethrow a process-control BaseException. - try: - self.agent.retain_handle_batch(handles) - except BaseException: - pass - while True: - try: - # Source-ownership barrier: only terminal, detached - # Works permit the scheduler to recycle this room. - self.agent.cancel_handles_and_wait(handles) - break - except BaseException as cleanup_error: - try: - if self.agent.handle_batch_is_quiescent(handles): - break - except BaseException: - pass - logger.warning( - "PyTorch cleanup for room %s was interrupted by %s; " - "retrying before publishing failure", - room, - type(cleanup_error).__qualname__, - ) - time.sleep(0) - if kv_chunk.staging_counted: - self._uncount_transfer_writer(room) - kv_chunk.staging_counted = False - if isinstance(e, Exception): - stored_error = e - try: - failure_message = str(e) - except BaseException: - failure_message = type(e).__qualname__ - else: - try: - detail = str(e) - except BaseException: - detail = "diagnostic formatting unavailable" - failure_message = f"{type(e).__qualname__}: {detail}" - # Do not retain a fatal exception traceback and its handle - # graph after every Work has become reusable. - stored_error = RuntimeError(failure_message) - self.exceptions[room] = stored_error - self.record_failure(room, failure_message) - self.update_status(room, KVPoll.Failed) - if self._transfer_writer_count( - room - ) == 0 and self.mark_transfer_quiescent(room): - self.transfer_infos.pop(room, None) - self.req_to_decode_prefix_len.pop(room, None) + except Exception as e: # Catch all exceptions to prevent silently killing this # worker thread, but still propagate via failure_exception(). if isinstance(e, _NIXL_TRANSPORT_ERRORS): @@ -1738,12 +1413,11 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) logger.exception( f"Unexpected transfer worker error for room {room}" ) - # Deferred-release acknowledgement now follows the terminal drain. - if not isinstance(e, Exception): - raise - finally: - if batch_active: - self.agent.end_handle_batch(handles) + self.exceptions[room] = e + self.record_failure(room, str(e)) + self.update_status(room, KVPoll.Failed) + # No ack here on purpose: the DONE barrier bails on the first + # ERR, so siblings may still be writing; fall back to the timeout. def register_buffer_to_engine(self): self.kv_descs = [] @@ -1811,94 +1485,12 @@ def _add_remote_peer(self, decode_kv_args: KVArgsRegisterInfo): decode_kv_args.requires_dcp_relayout = self.requires_dcp_relayout( decode_kv_args.dst_dcp_size, decode_kv_args.dst_dcp_rank ) - if self._use_torch_transfer: - if decode_kv_args.requires_dcp_relayout: - raise NotImplementedError( - "Experimental PyTorch transfer does not support DCP relayout" - ) - if decode_kv_args.decode_tp_size != self.attn_tp_size: - raise NotImplementedError( - "Experimental PyTorch transfer currently requires equal " - "prefill and decode attention TP sizes" - ) - if len(set(decode_kv_args.dst_kv_mem_kinds)) > 1: - raise NotImplementedError( - "Experimental PyTorch transfer does not support mixed remote " - "KV memory kinds" - ) if decode_kv_args.requires_dcp_relayout: self._init_dcp_pack_buffers_once(decode_kv_args.dst_dcp_size) - if self._use_torch_transfer: - missing = object() - prior_prepared = { - name: self.prep_handles.get(name, missing) for name in ("", agent_name) - } - try: - alias = self.agent.add_remote_agent( - decode_kv_args.agent_metadata, peer_name=agent_name - ) - if alias != agent_name: - raise RuntimeError( - "PyTorch transfer imported an unexpected peer alias " - f"{alias!r} for {agent_name!r}" - ) - if self.disaggregation_mode == DisaggregationMode.PREFILL: - self._prepare_payload_xfer(decode_kv_args) - # This table is the authoritative manager-level peer commit and - # therefore publishes last. - self.decode_kv_args_table[agent_name] = decode_kv_args - except BaseException as error: - if self.decode_kv_args_table.get(agent_name) is decode_kv_args: - self.decode_kv_args_table.pop(agent_name, None) - - # Quarantine derived prepared lists before peer rollback. The - # PyTorch path is equal-TP-only, so these are its only keys. - release_prepared = getattr(self.agent, "release_dlist_handle", None) - for name, previous in prior_prepared.items(): - current = self.prep_handles.get(name, missing) - if current is previous: - continue - if previous is missing: - self.prep_handles.pop(name, None) - else: - self.prep_handles[name] = previous - if current is not missing and callable(release_prepared): - try: - release_prepared(current) - except BaseException as cleanup_error: - add_note = getattr(error, "add_note", None) - if callable(add_note): - try: - add_note( - "prepared-list rollback also raised: " - f"{type(cleanup_error).__qualname__}" - ) - except BaseException: - pass - - # Exact metadata matching prevents a rejected duplicate from - # removing some other incarnation. TorchTransferAgent marks the - # matching peer retiring before its first provider close. - try: - rollback_peer = getattr(self.agent, "rollback_remote_agent", None) - if callable(rollback_peer): - rollback_peer(agent_name, decode_kv_args.agent_metadata) - except BaseException as cleanup_error: - add_note = getattr(error, "add_note", None) - if callable(add_note): - try: - add_note( - "PyTorch peer rollback also raised: " - f"{type(cleanup_error).__qualname__}" - ) - except BaseException: - pass - raise - else: - self.decode_kv_args_table[agent_name] = decode_kv_args - self.agent.add_remote_agent(decode_kv_args.agent_metadata) - if self.disaggregation_mode == DisaggregationMode.PREFILL: - self._prepare_payload_xfer(decode_kv_args) + self.decode_kv_args_table[agent_name] = decode_kv_args + self.agent.add_remote_agent(decode_kv_args.agent_metadata) + if self.disaggregation_mode == DisaggregationMode.PREFILL: + self._prepare_payload_xfer(decode_kv_args) def _send_kvcache_generic( self, @@ -2174,7 +1766,6 @@ def send_kvcache_mixed( prefill_kv_indices: npt.NDArray[np.int32], dst_kv_indices: npt.NDArray[np.int32], notif: str, - handles: list, ): info = self.decode_kv_args_table[peer_name] segments = info.kv_xfer_segments @@ -2183,6 +1774,7 @@ def send_kvcache_mixed( raise RuntimeError(f"Missing NIXL mixed KV transfer plan for {peer_name}") num_parts = len(segments) + handles = [] for part_idx, seg in enumerate(segments): num_layers = seg.end - seg.start src_indices = repeat_indices_over_layers( @@ -2205,8 +1797,8 @@ def send_kvcache_mixed( state = self.agent.transfer(xfer_handle) if state == "ERR": raise Exception("KVSender failed to post mixed prepped transfer") - if not self._use_torch_transfer: - handles.append(xfer_handle) + handles.append(xfer_handle) + return handles def send_kvcache_slice( self, @@ -2434,8 +2026,7 @@ def _do_staging_transfer( return (None, True) notif_tag = ( - f"{req.completion_notification_prefix}_stg_{kv_chunk.chunk_id}" - f"_{int(kv_chunk.is_last_chunk)}" + f"{req.room}_stg_{kv_chunk.chunk_id}_{int(kv_chunk.is_last_chunk)}" f"_{self.transfer_source_rank}_{chunk_idx}" f"_{page_start}_{num_pages}_{req.agent_name}" ) @@ -2755,8 +2346,6 @@ def maybe_send_extra( dst_state_item_lens: List[List[int]] | None = None, dst_state_dim_per_tensor: List[List[int]] | None = None, dst_state_layer_ids: List[List[int]] | None = None, - *, - handles: list, ): """Send state per hybrid component, dispatching by state_type[i].""" state_types = getattr(self.kv_args, "state_types", []) or [] @@ -2776,6 +2365,7 @@ def maybe_send_extra( dst_state_dim_per_tensor = dst_state_dim_per_tensor or [] dst_state_layer_ids = dst_state_layer_ids or [] + handles = [] for i, st in enumerate(state_types): src_indices = ( prefill_state_indices[i] if i < len(prefill_state_indices) else None @@ -2933,8 +2523,8 @@ def maybe_send_extra( f"PD Disaggregation via NIXL does NOT support {st} hybrid models yet." ) if h is not None: - if not self._use_torch_transfer: - handles.append(h) + handles.append(h) + return handles def add_transfer_request( self, @@ -2946,7 +2536,6 @@ def add_transfer_request( aux_index: Optional[int] = None, state_indices: Optional[List] = None, num_kv_tokens: Optional[int] = None, - wait_event: Optional[object] = None, ): assert self.disaggregation_mode == DisaggregationMode.PREFILL assert not is_last_chunk or (is_last_chunk and aux_index is not None) @@ -2958,11 +2547,7 @@ def add_transfer_request( if self.enable_staging: self._prefetch_staging_reqs(bootstrap_room) - if ( - bootstrap_room not in self.request_status - or self.check_status(bootstrap_room) == KVPoll.Failed - or bootstrap_room not in self.transfer_infos - ): + if bootstrap_room not in self.transfer_infos: # Dummy rank or already cleared; nothing to enqueue. return None @@ -2982,7 +2567,6 @@ def add_transfer_request( prefill_aux_index=aux_index, state_indices=state_indices, num_kv_tokens=num_kv_tokens, - wait_event=wait_event, ) ) return None @@ -2992,96 +2576,82 @@ def update_transfer_status(self): notif_map = self.agent.get_new_notifs() for peer_name, messages in notif_map.items(): for msg in messages: - try: - generation, room, tag, fields = _parse_completion_notification(msg) - except (TypeError, UnicodeError, ValueError) as exc: - logger.debug( - "Ignoring malformed NIXL completion from %s: %s", - peer_name, - exc, - ) - continue - - # This lookup never creates state. Generation comparison occurs - # before any status, staging-ring, or scatter mutation. - status = self.transfer_statuses.get(room) - if status is None or status.transfer_generation != generation: - logger.debug( - "Ignoring stale NIXL completion for room %s generation %s", - room, - generation, - ) - continue - - if tag in ("kv", "kvpart", "stg"): - source_rank = fields[2] - elif tag in ("aux", "state"): - source_rank = fields[0] - else: # The parser is exhaustive; retain a fail-closed guard. - continue - if not status.accepts_source_rank(source_rank): - logger.debug( - "Ignoring NIXL completion for unexpected source rank %s " - "in room %s generation %s", - source_rank, - room, - generation, - ) - continue - + # Notification tag layouts (underscore-separated): + # kv: {room}_kv_{chunk_id}_{is_last}_{pp_rank} -> 5 fields + # kvpart:{room}_kv_{chunk_id}_{is_last}_{pp_rank}_part_{i}_{n}-> 8 fields + # stg: {room}_stg_{chunk_id}_{is_last}_{pp_rank}_{chunk_idx} + # _{page_start}_{num_pages}_{agent_name} -> 9 fields + # aux: {room}_aux -> 2 fields + # state: {room}_state_{pp_rank} -> 3 fields + # maxsplit=8 keeps everything past the 8th underscore in the + # last component, so agent_name (which may itself contain + # underscores) lands intact in components[8] for the stg path. + components = msg.decode("ascii").split("_", 8) + room = int(components[0]) + tag = components[1] if tag == "kv": - self._track_kv_arrival(status, room, *fields) - elif tag == "kvpart": - self._track_kv_part_arrival(status, room, *fields) + chunk_id = int(components[2]) + is_last_chunk = bool(int(components[3])) + pp_rank = int(components[4]) if len(components) > 4 else 0 + if len(components) > 7 and components[5] == "part": + self._track_kv_part_arrival( + room, + chunk_id, + is_last_chunk, + pp_rank, + int(components[6]), + int(components[7]), + ) + else: + self._track_kv_arrival(room, chunk_id, is_last_chunk, pp_rank) elif tag == "stg": - self._handle_stg_notification(status, room, fields) + self._handle_stg_notification(components, room) elif tag == "aux": - self._handle_aux_notification(status, room, fields) + # Main's "nokv" marker carries the number of earlier KV + # chunks expected from this PP rank. + self._handle_aux_notification(room, components) elif tag == "state": - pp_rank, _component_id = fields - status.received_state_per_pp.add(pp_rank) + pp_rank = int(components[2]) if len(components) > 2 else 0 + self.transfer_statuses[room].received_state_per_pp.add(pp_rank) - def _handle_stg_notification( - self, status: TransferStatus, room: int, fields: Tuple[Any, ...] - ): + def _handle_stg_notification(self, components, room: int): """Handle a staging RDMA notification tag. - Format: nixlv2_{generation}_{room}_stg_{chunk_id}_{is_last}_{pp_rank}_{chunk_idx}_{page_start}_{num_pages}_{agent_name} + Format: {room}_stg_{chunk_id}_{is_last}_{pp_rank}_{chunk_idx}_{page_start}_{num_pages}_{agent_name} """ - ( - chunk_id, - is_last_chunk, - pp_rank, - chunk_idx, - page_start, - num_pages, - agent_name, - ) = fields + chunk_id = int(components[2]) + is_last_chunk = bool(int(components[3])) + pp_rank = int(components[4]) + chunk_idx = int(components[5]) + page_start = int(components[6]) + num_pages = int(components[7]) + agent_name = components[8] if len(components) > 8 else "" # Count this notif's own arrival BEFORE _track_kv_arrival, which can # conclude the transfer and record all-ranks Success. self._handle_staging_chunk_arrived( - room, chunk_idx, page_start, num_pages, pp_rank + room, chunk_idx, page_start, num_pages, agent_name ) - self._track_kv_arrival(status, room, chunk_id, is_last_chunk, pp_rank) + self._track_kv_arrival(room, chunk_id, is_last_chunk, pp_rank) - def _handle_aux_notification( - self, status: TransferStatus, room: int, fields: Tuple[Any, ...] - ): + def _handle_aux_notification(self, room: int, components: List[str]): """Handle an aux notification and trigger last scatter if staging is complete. Notification tag layouts: - aux: nixlv2_{generation}_{room}_aux_{pp_rank} - aux (nokv): nixlv2_{generation}_{room}_aux_nokv_{pp_rank}_{expected} + aux: {room}_aux -> 2 fields + aux (nokv): {room}_aux_nokv_{pp_rank}_{expected} -> 5 fields (the last chunk had no KV pages for this rank; `expected` is the number of prior KV chunks) """ - pp_rank, expected = fields - status.received_aux = True + self.transfer_statuses[room].received_aux = True # main's "nokv" marker (decode-side radix cache hit, see #19746). - if expected is not None: - status.expected_kvs_per_pp[pp_rank] = expected - if status.num_pp_ranks_expected is None: - status.num_pp_ranks_expected = self.required_prefill_response_num_table.get( + if len(components) > 3 and components[2] == "nokv": + pp_rank = int(components[3]) + expected = int(components[4]) if len(components) > 4 else 0 + self.transfer_statuses[room].expected_kvs_per_pp[pp_rank] = expected + if self.transfer_statuses[room].num_pp_ranks_expected is None: + self.transfer_statuses[ + room + ].num_pp_ranks_expected = self.required_prefill_response_num_table.get( room, 1 ) if ( @@ -3089,34 +2659,30 @@ def _handle_aux_notification( and self._staging_handler is not None and self._staging_handler.is_staging_room(room) ): - self._maybe_submit_last_scatter(room, status) + self._maybe_submit_last_scatter(room) def _track_kv_arrival( - self, - status: TransferStatus, - room: int, - chunk_id: int, - is_last_chunk: bool, - pp_rank: int, + self, room: int, chunk_id: int, is_last_chunk: bool, pp_rank: int ): """Update transfer status tracking for a kv chunk arrival.""" - status.received_kvs_per_pp[pp_rank].add(chunk_id) + self.transfer_statuses[room].received_kvs_per_pp[pp_rank].add(chunk_id) if is_last_chunk: - status.expected_kvs_per_pp[pp_rank] = chunk_id + 1 - if status.num_pp_ranks_expected is None: - status.num_pp_ranks_expected = ( - self.required_prefill_response_num_table.get(room, 1) + self.transfer_statuses[room].expected_kvs_per_pp[pp_rank] = chunk_id + 1 + if self.transfer_statuses[room].num_pp_ranks_expected is None: + self.transfer_statuses[ + room + ].num_pp_ranks_expected = self.required_prefill_response_num_table.get( + room, 1 ) if ( self.enable_staging and self._staging_handler is not None and self._staging_handler.is_staging_room(room) ): - self._maybe_submit_last_scatter(room, status) + self._maybe_submit_last_scatter(room) def _track_kv_part_arrival( self, - status: TransferStatus, room: int, chunk_id: int, is_last_chunk: bool, @@ -3126,7 +2692,7 @@ def _track_kv_part_arrival( ): """Track one segment of a mixed-memory KV transfer.""" if num_parts <= 1: - self._track_kv_arrival(status, room, chunk_id, is_last_chunk, pp_rank) + self._track_kv_arrival(room, chunk_id, is_last_chunk, pp_rank) return if part_idx < 0 or part_idx >= num_parts: raise RuntimeError( @@ -3136,6 +2702,7 @@ def _track_kv_part_arrival( ) key = (pp_rank, chunk_id) + status = self.transfer_statuses[room] if status.received_kv_parts_per_pp is None: status.received_kv_parts_per_pp = defaultdict(set) if status.expected_kv_parts_per_pp is None: @@ -3151,7 +2718,7 @@ def _track_kv_part_arrival( if len(parts) == num_parts: status.received_kv_parts_per_pp.pop(key, None) status.expected_kv_parts_per_pp.pop(key, None) - self._track_kv_arrival(status, room, chunk_id, is_last_chunk, pp_rank) + self._track_kv_arrival(room, chunk_id, is_last_chunk, pp_rank) def _handle_staging_chunk_arrived( self, @@ -3159,7 +2726,7 @@ def _handle_staging_chunk_arrived( chunk_idx: int, page_start: int, num_pages: int, - source_rank: int, + agent_name: str, ): """Process a staging chunk arrival via RDMA notification.""" handler = self._staging_handler @@ -3170,19 +2737,14 @@ def _handle_staging_chunk_arrived( chunk_idx, page_start, num_pages, - source_rank, + agent_name, ) - def _maybe_submit_last_scatter( - self, room: int, status: Optional[TransferStatus] = None - ): + def _maybe_submit_last_scatter(self, room: int): """Check if all kv+aux transfers are done and submit last scatter if so.""" - current_status = self.transfer_statuses.get(room) - if current_status is None or ( - status is not None and current_status is not status - ): + status = self.transfer_statuses.get(room) + if status is None: return - status = current_status if not status.received_aux: return if status.num_pp_ranks_expected is None: @@ -3206,22 +2768,15 @@ def _handle_abort_notification(self, msg: List[bytes]) -> bool: return False try: - if len(msg) != 5: - raise ValueError("ABORT must contain exactly five frames") room_to_be_aborted = int(msg[1].decode("ascii")) - decode_ip = msg[2].decode("ascii") - decode_port = int(msg[3].decode("ascii")) - transfer_generation = validate_transfer_generation(msg[4].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}") return True - generation_active = ( - self.active_transfer_generation(room_to_be_aborted) == transfer_generation - ) room_active = ( - generation_active - and room_to_be_aborted in self.request_status + room_to_be_aborted in self.request_status and self.check_status(room_to_be_aborted) != KVPoll.Success ) if room_active: @@ -3246,168 +2801,78 @@ def _handle_abort_notification(self, msg: List[bytes]) -> bool: # 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 decode_port is not None and self.register_deferred_ack_target( - room_to_be_aborted, - transfer_generation, - decode_ip, - decode_port, - ): - self._maybe_ack_drained_abort(room_to_be_aborted, transfer_generation) - elif not generation_active: - logger.debug( - "Ignoring stale ABORT for room %s generation %s", - room_to_be_aborted, - transfer_generation, - ) + 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) return True - def _handle_bootstrap_message(self, waiting_req_bytes): - """Handle one bootstrap message and persist failures for propagation.""" - - agent_name: Optional[str] = None - room_id: Optional[int] = None - is_room_message = False - room_identity_authenticated = False - try: - logger.debug( - "Received multipart with total byte size %d", - sum(len(x) for x in waiting_req_bytes), - ) - if not waiting_req_bytes: - raise ValueError("NIXL bootstrap message is empty") - - # Staging: decode reports consumption watermark back to prefill. - if waiting_req_bytes[0] == b"WATERMARK": - if self.enable_staging: - handle_watermark_msg(self._staging_ctx, waiting_req_bytes) - return + def _start_bootstrap_thread(self): + def bootstrap_thread(): + """This thread recvs transfer info from the decode engine""" + while True: + waiting_req_bytes = self.server_socket.recv_multipart() + logger.debug( + f"Received multipart with total byte size {sum(len(x) for x in waiting_req_bytes)}" + ) - # Staging: decode replies with allocated staging offset. - if waiting_req_bytes[0] == STAGING_RSP_WIRE_TAG: - if self.enable_staging: - handle_staging_rsp(waiting_req_bytes, self.transfer_infos) - return + # Staging: decode reports consumption watermark back to prefill + if waiting_req_bytes[0] == b"WATERMARK": + if self.enable_staging: + handle_watermark_msg(self._staging_ctx, waiting_req_bytes) + continue - if self._handle_abort_notification(waiting_req_bytes): - return + # Staging: decode replies with allocated staging offset + if waiting_req_bytes[0] == b"STAGING_RSP": + if self.enable_staging: + handle_staging_rsp(waiting_req_bytes, self.transfer_infos) + continue - if waiting_req_bytes[0] != GUARD: - raise ValueError(f"First message should be {GUARD}. Foreign traffic?") - payload = waiting_req_bytes[1:] - room_text = payload[0].decode("ascii") - agent_name = payload[3].decode("ascii") - if room_text == "None": - # Registration failures are agent-scoped because no room is - # available yet. A subsequent room message propagates them. - self._add_remote_peer(KVArgsRegisterInfo.from_zmq(payload)) - self._bootstrap_errors.pop(agent_name, None) - logger.debug("Register KVArgs from %s successfully", agent_name) - return + if self._handle_abort_notification(waiting_req_bytes): + continue - is_room_message = True - room_id = int(room_text) - try: - incoming_generation = validate_transfer_generation( - payload[10].decode("ascii") if len(payload) > 10 else None - ) - except (AttributeError, UnicodeError, ValueError) as error: - # A room number without an exact generation cannot own or fail - # that recyclable room, nor poison this agent's later traffic. - logger.warning( - "Dropping unauthenticated NIXL metadata for room %s: %s", - room_id, - error, - ) - return - prior_error = self._bootstrap_errors.get(agent_name) - if prior_error is not None: - # Canonical generation plus the failed peer registration is the - # ownership evidence; no active-room table was necessarily - # initialized when peer setup failed. - room_identity_authenticated = True - raise prior_error - active_generation = self.active_transfer_generation(room_id) - if ( - active_generation is not None - and active_generation != incoming_generation - ): - logger.warning( - "Dropping late NIXL metadata for room %s generation %s; " - "active generation is %s", - room_id, - incoming_generation, - active_generation, + assert waiting_req_bytes[0] == GUARD, ( + f"First message should be {GUARD}. Foreign traffic?" ) - return - room_identity_authenticated = active_generation == incoming_generation - transfer_info = TransferInfo.from_zmq(payload) - room_identity_authenticated = True - if not self.register_transfer_generation( - room_id, transfer_info.transfer_generation - ): - logger.warning( - "Dropping late NIXL metadata for room %s generation %s; " - "active generation is %s", - room_id, - transfer_info.transfer_generation, - self.active_transfer_generation(room_id), - ) - return - if room_id not in self.transfer_infos: - self.transfer_infos[room_id] = {} - self.transfer_infos[room_id][agent_name] = transfer_info - required_dst_info_num = self.transfer_infos[room_id][ - agent_name - ].required_dst_info_num - logger.debug( - "got info room=%s agent_name=%s required_dst_info_num=%s", - room_id, - agent_name, - required_dst_info_num, - ) - if len(self.transfer_infos[room_id]) == required_dst_info_num: - self.resolve_kv_replica_factor(self.transfer_infos[room_id]) - self.req_to_decode_prefix_len[room_id] = next( - ( - info.decode_prefix_len - for info in self.transfer_infos[room_id].values() - if info.decode_prefix_len is not None - ), - 0, - ) - logger.debug("room=%s is bootstrapped", room_id) - self.update_status(room_id, KVPoll.WaitingForInput) - except Exception as error: - if is_room_message and not room_identity_authenticated: - logger.warning( - "Dropping malformed unauthenticated NIXL room metadata: %s", - error, + waiting_req_bytes = waiting_req_bytes[1:] + room = waiting_req_bytes[0].decode("ascii") + agent_name = waiting_req_bytes[3].decode("ascii") + if room == "None": + # Register new peer and save KV base pointers. + self._add_remote_peer( + KVArgsRegisterInfo.from_zmq(waiting_req_bytes) + ) + logger.debug(f"Register KVArgs from {agent_name} successfully") + continue + room = int(room) + if room not in self.transfer_infos: + self.transfer_infos[room] = {} + self.transfer_infos[room][agent_name] = TransferInfo.from_zmq( + waiting_req_bytes ) - return - if agent_name is not None: - self._bootstrap_errors.setdefault(agent_name, error) - error = self._bootstrap_errors[agent_name] - if room_id is not None: - self.exceptions[room_id] = error - self.record_failure(room_id, str(error)) - self.update_status(room_id, KVPoll.Failed) - raise - - def _start_bootstrap_thread(self): - - def bootstrap_thread(): - """This thread recvs transfer info from the decode engine""" - while True: - waiting_req_bytes = self.server_socket.recv_multipart() - try: - self._handle_bootstrap_message(waiting_req_bytes) - except Exception: - logger.exception( - "NIXL bootstrap message failed; listener remains active" + required_dst_info_num = self.transfer_infos[room][ + agent_name + ].required_dst_info_num + logger.debug(f"got info {room=} {agent_name=} {required_dst_info_num=}") + if len(self.transfer_infos[room]) == required_dst_info_num: + self.resolve_kv_replica_factor(self.transfer_infos[room]) + self.req_to_decode_prefix_len[room] = next( + ( + info.decode_prefix_len + for info in self.transfer_infos[room].values() + if info.decode_prefix_len is not None + ), + 0, ) + logger.debug(f"{room=} is bootstrapped") + self.update_status(room, KVPoll.WaitingForInput) - threading.Thread(target=bootstrap_thread, daemon=True).start() + threading.Thread(target=bootstrap_thread).start() class NixlKVSender(CommonKVSender): @@ -3434,16 +2899,6 @@ def __init__( self._send_failed = False self._send_error: Optional[Exception] = None self._transfer_start_time: Optional[float] = None - self._early_send_wait_event: Optional[object] = None - - @property - def supports_overlap_early_send(self) -> bool: - """Whether the data plane can order RDMA after a CUDA readiness event.""" - - # Both native NIXL and the PyTorch transfer provider use the common - # transfer worker, which consumes the readiness event before touching - # source memory. - return True def send( self, @@ -3465,9 +2920,6 @@ def send( ): self._transfer_start_time = time.perf_counter() - wait_event = getattr(self, "_early_send_wait_event", None) - self._early_send_wait_event = None - self.kv_mgr.add_transfer_request( self.bootstrap_room, kv_indices, @@ -3477,7 +2929,6 @@ def send( self.aux_index, state_indices, num_kv_tokens, - wait_event=wait_event, ) self._record_transfer_indices(kv_indices, state_indices) self.chunk_id += 1 @@ -3547,51 +2998,11 @@ def __init__( mgr: NixlKVManager, bootstrap_addr: str, bootstrap_room: Optional[int] = None, - transfer_generation: Optional[str] = None, ): self.started_transfer = False - super().__init__( - mgr, - bootstrap_addr, - bootstrap_room, - transfer_generation=transfer_generation, - ) - # Arm the exact expected generation before metadata can cause any - # completion. Reuse replaces old room status; late old-generation - # notifications are rejected by update_transfer_status(). - self.kv_mgr.transfer_statuses[self.bootstrap_room] = TransferStatus( - transfer_generation=self.transfer_generation - ) + super().__init__(mgr, bootstrap_addr, bootstrap_room) self.init_time = None - def init(self, prefill_dp_rank: int): - super().init(prefill_dp_rank) - if self.conclude_state == KVPoll.Failed: - return - expected_source_ranks = frozenset( - bootstrap_info["prefill_unique_rank"] - for bootstrap_info in self.bootstrap_infos - if not bootstrap_info["is_dummy"] - ) - status = self.kv_mgr.transfer_statuses.get(self.bootstrap_room) - if status is None or status.transfer_generation != self.transfer_generation: - raise RuntimeError( - "NIXL completion status lost its receiver lease identity" - ) - status.bind_expected_source_ranks(expected_source_ranks) - - def _clear_own_transfer_status(self) -> None: - status = self.kv_mgr.transfer_statuses.get(self.bootstrap_room) - if ( - status is not None - and status.transfer_generation == self.transfer_generation - ): - self.kv_mgr.transfer_statuses.pop(self.bootstrap_room, None) - - def clear(self) -> None: - self._clear_own_transfer_status() - super().clear() - def send_metadata( self, kv_indices: npt.NDArray[np.int32], @@ -3647,7 +3058,6 @@ def send_metadata( packed_state_indices, str(decode_prefix_len or 0).encode("ascii"), str(int(is_dummy)).encode("ascii"), - self.transfer_generation.encode("ascii"), ] ) except zmq.ZMQError: @@ -3680,18 +3090,18 @@ def poll(self) -> KVPoll: if not self.started_transfer: return status - # Drain notifications before enforcing the waiting deadline. Both roles - # use transport background progress, but application notifications are - # ingested into room state only here via update_transfer_status(); a - # completion queued by NIXL at/after the deadline would otherwise lose - # to the timeout purely by poll ordering. + # Drain notifications before enforcing the waiting deadline. The decode + # agent has no NIXL progress thread (num_threads=0), so incoming + # completion notifications are only ingested here via + # update_transfer_status(); a completion queued by NIXL at/after the + # deadline would otherwise lose to the timeout purely by poll ordering. self.kv_mgr.update_transfer_status() if self.kv_mgr.check_transfer_done(self.bootstrap_room): # type: ignore self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].discard( self.bootstrap_room ) self.conclude_state = KVPoll.Success - self._clear_own_transfer_status() + del self.kv_mgr.transfer_statuses[self.bootstrap_room] return self.conclude_state # type: ignore timeout_result = self._check_waiting_timeout() diff --git a/python/sglang/srt/disaggregation/nixl/torch_transfer.py b/python/sglang/srt/disaggregation/nixl/torch_transfer.py deleted file mode 100644 index 7693b7ba2e97..000000000000 --- a/python/sglang/srt/disaggregation/nixl/torch_transfer.py +++ /dev/null @@ -1,2169 +0,0 @@ -"""Experimental PyTorch endpoint-transfer adapter for SGLang NIXL PD.""" - -# Process-control exceptions are deliberately part of the ownership protocol: -# every BaseException boundary below reconciles or durably retains native state. -# ruff: noqa: BLE001 - -from __future__ import annotations - -import importlib -import itertools -import logging -import operator -import threading -import time -import uuid -from collections import defaultdict -from collections.abc import Callable, Mapping, Sequence -from contextlib import contextmanager -from dataclasses import dataclass, field -from typing import Any, Literal - -logger = logging.getLogger(__name__) - -_UNKNOWN_NOTIFICATION_SOURCE = "" -_RETAINED_WORK_SWEEP_INTERVAL_S = 0.01 -_PENDING_STATES = frozenset({"CREATED", "PENDING", "RUNNING", "SUBMITTED"}) -_DONE_STATES = frozenset({"COMPLETE", "COMPLETED", "DONE", "SUCCEEDED"}) -_FAILED_STATES = frozenset({"CANCELLED", "FAILED", "TIMED_OUT"}) -_UNKNOWN_WORK_OUTCOME = object() -_FACTORY_API_VERSION = 2 - - -def _load_transfer_api() -> Any: - try: - module = importlib.import_module("torch.distributed._transfer") - except ImportError as exc: - raise ImportError( - "PyTorch endpoint transfer requires torch.distributed._transfer" - ) from exc - factory_version = getattr(module, "BACKEND_FACTORY_API_VERSION", None) - if ( - type(factory_version) is not int - or factory_version < _FACTORY_API_VERSION - or not hasattr(module, "BackendFactoryV2") - or not hasattr(module, "Endpoint") - or not hasattr(module, "RawSpan") - or not hasattr(module, "TransferOp") - or not callable( - getattr( - getattr(module, "TransferPlan", None), - "submit_indices_and_poll", - None, - ) - ) - ): - raise ImportError( - "PyTorch endpoint transfer requires FactoryV2 Core, Endpoint, RawSpan, " - "TransferOp, and allocation-minimal exact-state indexed submission " - "support" - ) - return module - - -def _register_nixl_backend() -> None: - """Load the external provider only after explicit SGLang opt-in.""" - try: - module = importlib.import_module("nixl.torch_transfer") - except ImportError as exc: - raise ImportError( - "The PyTorch NIXL transfer backend requires nixl.torch_transfer" - ) from exc - register = getattr(module, "register_torch_backend", None) - factory_version = getattr(module, "TORCH_TRANSFER_FACTORY_API_VERSION", None) - if type(factory_version) is not int or factory_version != _FACTORY_API_VERSION: - raise ImportError( - "nixl.torch_transfer must register the exact FactoryV2 provider ABI" - ) - if not callable(register): - raise ImportError( # noqa: TRY004 - "nixl.torch_transfer does not expose register_torch_backend()" - ) - register() - - -@dataclass(frozen=True) -class _RawDescriptor: - address: int - nbytes: int - device_id: int - stride: int | None = None - count: int | None = None - - @property - def extent(self) -> int: - if self.stride is None or self.count is None: - return self.nbytes - return (self.count - 1) * self.stride + self.nbytes - - -@dataclass(frozen=True) -class _RawDescriptorList: - descriptors: tuple[_RawDescriptor, ...] - memory_type: str - - -@dataclass(frozen=True) -class _PreparedList: - regions: tuple[object, ...] - remote: bool - peer_name: str | None = None - - -@dataclass(frozen=True) -class _CachedPlan: - plan: object - peer_name: str - prepared_ids: frozenset[int] - - -@dataclass -class _RegistrationRecord: - address: int - nbytes: int - device_id: int - memory_type: str - registration: object - - -def _freeze_index_vector(indices: Sequence[int]) -> Sequence[int]: - """Take immutable ownership without copying SGLang's native-int32 fast path. - - SGLang creates fresh, owning NumPy int32 arrays immediately before - ``make_prepped_xfer``. Mark those arrays read-only and retain them until - Core snapshots them during the adjacent ``transfer`` call. Non-owning - buffers are copied before freezing, and general Python sequences become an - immutable normalized tuple. Thus no caller-visible mutation can alter an - admitted request while the production fast path adds no index memcpy. - """ - - try: - view = memoryview(indices) - except TypeError: - return tuple(operator.index(value) for value in indices) - try: - native_int32 = ( - view.ndim == 1 - and view.c_contiguous - and view.itemsize == 4 - and view.format in {"i", "@i", "=i"} - ) - if not native_int32: - return tuple(operator.index(value) for value in indices) - if view.readonly: - return indices - flags = getattr(indices, "flags", None) - setflags = getattr(indices, "setflags", None) - if ( - flags is not None - and bool(getattr(flags, "owndata", False)) - and callable(setflags) - ): - setflags(write=False) - return indices - copy = getattr(indices, "copy", None) - if callable(copy): - frozen = copy() - freeze = getattr(frozen, "setflags", None) - if callable(freeze): - freeze(write=False) - return frozen - return tuple(operator.index(value) for value in indices) - finally: - view.release() - - -@dataclass(eq=False) -class _WorkHandle: - work: object | None = None - owned_plan: object | None = None - terminal_state: str | None = None - first_state: object | None = None - operation: Literal["READ", "WRITE"] | None = None - plan: object | None = None - local_indices: Sequence[int] | None = None - remote_indices: Sequence[int] | None = None - notification: bytes | None = None - peer_name: str | None = None - cache_key: tuple[str, int, int] | None = None - # The production token is the already allocated handle itself. This avoids - # a second per-attempt allocation; Core Work.close breaks the temporary - # Work-to-handle edge before the adapter drops its registry ownership. - submission_interrupted: bool = False - _state_lock: threading.Lock = field( - default_factory=threading.Lock, init=False, repr=False - ) - _status_in_progress: bool = field(default=False, init=False, repr=False) - _recovery_in_progress: bool = field(default=False, init=False, repr=False) - _release_in_progress: bool = field(default=False, init=False, repr=False) - _native_submit_entered: bool = field(default=False, init=False, repr=False) - phase: Literal[ - "prepared", - "posting", - "recovery_pending", - "posted", - "terminal", - "released", - ] = "prepared" - - -class TorchTransferAgent: - """NIXL-agent-shaped boundary around ``torch.distributed._transfer``. - - It exists to validate the current SGLang data plane without spreading an - experimental PyTorch API through ``conn.py``. Core ``MULTIPLE`` mode lets - independent transfer-worker calls overlap while a short lifecycle gate - drains every admitted call before close. Submission - deliberately uses Core's ``caller_ready`` mode: the common transfer worker - consumes the early-send CUDA event before staging, packing, or direct source - access, and it waits for terminal Work before any chunk is reusable. - """ - - def __init__( - self, - name: str, - *, - owner: object, - background_progress: bool, - num_threads: int = 0, - endpoint_backend: str = "nixl", - ) -> None: - self.name = name - self._owner = owner - self._background_progress = background_progress - self._num_threads = num_threads - self._endpoint_backend = endpoint_backend - self._api = _load_transfer_api() - self._core_transfer_operations = { - "READ": self._api.TransferOp.READ, - "WRITE": self._api.TransferOp.WRITE, - } - self._work_state_outcomes = self._complete_work_state_outcomes(self._api) - self._lock = threading.RLock() - self._lifecycle_cv = threading.Condition(self._lock) - self._admission_local = threading.local() - self._active_calls = 0 - self._closing = False - self._close_owner: int | None = None - self._endpoint: Any = None - self._endpoint_recovery: Any = None - self._transport_backend: str | None = None - self._send_notification_sync: Callable[[object, bytes], None] | None = None - self._registrations: list[_RegistrationRecord] = [] - self._registration_groups: list[tuple[object, ...]] = [] - self._registration_ids = itertools.count() - self._peers: dict[str, object] = {} - self._peer_wire_metadata: dict[str, bytes] = {} - self._source_aliases: dict[object, str] = {} - self._plan_cache: dict[tuple[str, int, int], _CachedPlan] = {} - self._closing_plans: dict[tuple[str, int, int], _CachedPlan] = {} - self._retiring_peers: set[str] = set() - self._works: dict[int, _WorkHandle] = {} - self._batch_handles = threading.local() - self._active_handle_batches: dict[int, list[_WorkHandle]] = {} - self._retained_handle_batches: dict[int, list[_WorkHandle]] = {} - self._detached_works: list[_WorkHandle] = [] - self._retained_work_ids: set[int] = set() - self._sweep_wakeup = threading.Event() - self._sweep_thread: threading.Thread | None = None - # Set before entering any Core child constructor and cleared only after - # authoritative adapter publication or a complete identity reconcile. - self._construction_dirty = False - - def create_backend(self, backend: str, backend_params: Mapping[str, str]) -> None: - with self._lock: - self._retry_endpoint_cleanup_unlocked() - if self._endpoint is not None: - raise RuntimeError("PyTorch transfer endpoint is already initialized") - options = { - "backends": [backend], - "backend_init_params": {backend: dict(backend_params)}, - "transfer_backends": [backend], - "num_threads": self._num_threads, - } - if self._endpoint_backend == "nixl": - _register_nixl_backend() - endpoint_id = str(uuid.uuid4()) - progress_mode = ( - self._api.ProgressMode.BACKGROUND - if self._background_progress - else self._api.ProgressMode.MANUAL - ) - thread_mode = self._api.ThreadMode.MULTIPLE - endpoint = None - try: - endpoint = self._api.Endpoint( - self.name, - backend=self._endpoint_backend, - endpoint_id=endpoint_id, - progress_mode=progress_mode, - # Core/provider synchronization owns native concurrency. - # Adapter locks protect only short Python state transitions. - thread_mode=thread_mode, - options=options, - ) - required_capabilities = ( - "raw_spans", - "strided_regions", - "write", - "attached_notifications", - "standalone_notifications", - ) - missing = [ - name - for name in required_capabilities - if not self._supports_capability(endpoint, name) - ] - if missing: - raise NotImplementedError( - "The selected PyTorch transfer backend is missing required " - f"capabilities: {', '.join(missing)}" - ) - send_notification_sync = getattr( - endpoint, "send_notification_sync", None - ) - if not ( - self._supports_capability(endpoint, "synchronous_notification_send") - and callable(send_notification_sync) - ): - send_notification_sync = None - # Derived publication precedes the authoritative Endpoint owner. - self._send_notification_sync = send_notification_sync - self._transport_backend = backend - self._endpoint = endpoint - except BaseException as error: - # A process-control exception may land after the authoritative - # assignment but before this method returns. The failed call - # still owns rollback: never leave a usable Endpoint published. - if endpoint is not None and self._endpoint is endpoint: - self._endpoint = None - self._send_notification_sync = None - self._transport_backend = None - - candidate = endpoint - if candidate is None: - try: - recovery_endpoint = getattr(error, "recovery_endpoint", None) - except BaseException as lookup_error: - recovery_endpoint = None - self._add_exception_note( - error, - "reading recovery_endpoint also raised: " - f"{type(lookup_error).__qualname__}", - ) - if recovery_endpoint is not None: - try: - trusted = self._matches_recovery_endpoint( - recovery_endpoint, - endpoint_id=endpoint_id, - progress_mode=progress_mode, - thread_mode=thread_mode, - ) - except BaseException as match_error: - trusted = False - self._add_exception_note( - error, - "validating recovery_endpoint also raised: " - f"{type(match_error).__qualname__}", - ) - if trusted: - candidate = recovery_endpoint - else: - # Exact identity fields are the ownership proof. - # Even another real Core Endpoint can belong to an - # unrelated live caller and must not be mutated. - self._add_exception_note( - error, - "ignored a non-matching recovery_endpoint during " - "SGLang endpoint construction", - ) - if candidate is not None: - self._endpoint_recovery = candidate - try: - self._retry_endpoint_cleanup_unlocked() - except BaseException as cleanup_error: - self._add_exception_note( - error, - "endpoint construction cleanup also raised: " - f"{type(cleanup_error).__qualname__}", - ) - raise - - def get_plugin_list(self) -> list[str]: - with self._lock: - return [self._transport_backend] if self._transport_backend else [] - - def register_memory( - self, descriptors: Sequence[Sequence[object]], memory_type: str - ) -> tuple[object, ...]: - with self._lock: - self._reconcile_construction_state_unlocked() - endpoint = self._require_endpoint() - self._construction_dirty = True - registrations = [] - records = [] - group = None - try: - for raw in self._normalize_descriptors(descriptors): - span = self._api.RawSpan( - address=raw.address, - nbytes=raw.extent, - device=self._device(memory_type, raw.device_id), - memory_type=memory_type, - owner=self._owner, - ) - name = f"sglang_region_{next(self._registration_ids)}" - registration = endpoint.register(span, name=name) - registrations.append(registration) - records.append( - _RegistrationRecord( - address=raw.address, - nbytes=raw.extent, - device_id=raw.device_id, - memory_type=memory_type, - registration=registration, - ) - ) - group = tuple(registrations) - # The complete registration graph is authoritative only after - # both indexes publish. Exception recovery removes either partial. - self._registration_groups.append(group) - self._registrations.extend(records) - self._construction_dirty = False - return group - except BaseException as error: - self._registrations = [ - record - for record in self._registrations - if not any(record is candidate for candidate in records) - ] - if group is not None: - self._registration_groups = [ - candidate - for candidate in self._registration_groups - if candidate is not group - ] - self._reconcile_after_construction_error_unlocked(error, "registration") - raise - - def deregister_memory(self, registrations: Sequence[object]) -> None: - with self._lock: - self._require_endpoint() - self._reconcile_construction_state_unlocked() - self._require_no_unsettled_interrupted_submission("deregister memory") - for registration in reversed(tuple(registrations)): - registration.close() - removed = set(map(id, registrations)) - self._registrations = [ - record - for record in self._registrations - if id(record.registration) not in removed - ] - self._registration_groups = [ - remaining - for group in self._registration_groups - if ( - remaining := tuple( - registration - for registration in group - if id(registration) not in removed - ) - ) - ] - - def get_agent_metadata(self) -> bytes: - with self._lock: - endpoint = self._require_endpoint() - metadata = endpoint.export_metadata( - [record.registration for record in self._registrations] - ) - if isinstance(metadata, bytes): - return metadata - to_bytes = getattr(metadata, "to_bytes", None) - if callable(to_bytes): - return to_bytes() - try: - return bytes(metadata) - except (TypeError, ValueError) as exc: - raise TypeError("Endpoint metadata is not bytes-compatible") from exc - - def add_remote_agent(self, metadata: bytes, peer_name: str | None = None) -> str: - with self._lock: - self._reconcile_construction_state_unlocked() - metadata = bytes(metadata) - if peer_name is not None and peer_name in self._retiring_peers: - self.remove_remote_agent(peer_name) - if peer_name is not None and peer_name in self._peers: - if self._peer_wire_metadata.get(peer_name) == metadata: - return peer_name - raise ValueError( - f"Peer {peer_name!r} is already registered with different metadata" - ) - import_options = {} if peer_name is None else {"expected_name": peer_name} - self._construction_dirty = True - peer = None - alias = None - aliases = None - previous_aliases = self._source_aliases - try: - peer = self._require_endpoint().import_peer(metadata, **import_options) - source_ids = self._peer_source_ids(peer) - source_id = source_ids[0] - alias = peer_name or source_id - if alias in self._peers: - peer.close() - raise ValueError(f"Peer {alias!r} is already registered") - if not callable(getattr(peer, "unsafe_region", None)): - peer.close() - raise NotImplementedError( - "Imported peer does not support provider-validated raw regions" - ) - aliases = dict(previous_aliases) - for source in source_ids: - aliases[source] = alias - incarnation = getattr(peer, "incarnation", None) - if incarnation is not None: - aliases[(source_id, str(incarnation))] = alias - self._source_aliases = aliases - self._peer_wire_metadata[alias] = metadata - # The exact Core peer is the authoritative adapter owner and - # commits last, after every derived identity index. - self._peers[alias] = peer - self._construction_dirty = False - return alias - except BaseException as error: - if alias is None or self._peers.get(alias) is not peer: - if aliases is not None and self._source_aliases is aliases: - self._source_aliases = previous_aliases - if alias is not None: - self._peer_wire_metadata.pop(alias, None) - self._reconcile_after_construction_error_unlocked(error, "peer") - else: - # A dict STORE may commit and then raise. Exact map identity - # proves adoption, so retry is idempotent and no scan is due. - self._construction_dirty = False - raise - - def rollback_remote_agent(self, peer_name: str, metadata: bytes) -> bool: - """Remove only the peer carrying this exact bootstrap wire metadata.""" - - with self._lock: - self._reconcile_construction_state_unlocked() - if self._peer_wire_metadata.get(peer_name) != bytes(metadata): - return False - self.remove_remote_agent(peer_name) - return True - - def remove_remote_agent(self, peer_name: str) -> None: - with self._lock: - self._require_endpoint() - self._reconcile_construction_state_unlocked() - peer = self._peers.get(peer_name) - if peer is None: - self._retiring_peers.discard(peer_name) - return - # Quarantine before the first provider close. If close partially - # commits and raises, no future submission may reuse this peer. - self._retiring_peers.add(peer_name) - self._require_no_unsettled_interrupted_submission("remove a remote agent") - self._release_cached_plans(lambda cached: cached.peer_name == peer_name) - if not getattr(peer, "closed", False): - peer.close() - self._source_aliases = { - source: alias - for source, alias in self._source_aliases.items() - if alias != peer_name - } - self._peer_wire_metadata.pop(peer_name, None) - self._peers.pop(peer_name, None) - self._retiring_peers.discard(peer_name) - - def get_xfer_descs( - self, descriptors: Sequence[Sequence[object]], memory_type: str - ) -> _RawDescriptorList: - return _RawDescriptorList(self._normalize_descriptors(descriptors), memory_type) - - def prep_xfer_dlist( - self, - peer_name: str, - descriptors: Sequence[Sequence[object]], - memory_type: str, - ) -> _PreparedList: - with self._lock: - self._require_endpoint() - raw = self._normalize_descriptors(descriptors) - if peer_name: - regions = self._resolve_remote(peer_name, raw, memory_type) - return _PreparedList(regions, remote=True, peer_name=peer_name) - return _PreparedList(self._resolve_local(raw, memory_type), remote=False) - - def release_dlist_handle(self, handle: _PreparedList) -> None: - with self._lock: - self._require_endpoint() - prepared_id = id(handle) - self._release_cached_plans( - lambda cached: prepared_id in cached.prepared_ids - ) - - def make_prepped_xfer( - self, - operation: str, - source: _PreparedList, - source_indices: Sequence[int], - destination: _PreparedList, - destination_indices: Sequence[int], - notification: bytes = b"", - ) -> _WorkHandle: - with self._admitted_call(): - operation = self._normalize_operation(operation) - local, remote, local_indices, remote_indices = self._orient_prepared( - operation, - source, - source_indices, - destination, - destination_indices, - ) - local_indices = _freeze_index_vector(local_indices) - remote_indices = _freeze_index_vector(remote_indices) - key = (operation, id(local), id(remote)) - wire_notification = bytes(notification) if notification else None - handle = None - with self._lock: - if remote.peer_name in self._retiring_peers: - raise RuntimeError( - f"Peer {remote.peer_name!r} teardown is in progress" - ) - if key in self._closing_plans: - raise RuntimeError("Prepared transfer plan teardown is in progress") - cached = self._plan_cache.get(key) - if cached is None: - assert remote.peer_name is not None - self._reconcile_construction_state_unlocked() - self._construction_dirty = True - cached = None - try: - plan = self._require_endpoint().prepare( - local=local.regions, - remote=remote.regions, - indexed=True, - ) - cached = _CachedPlan( - plan=plan, - peer_name=remote.peer_name, - prepared_ids=frozenset({id(local), id(remote)}), - ) - self._plan_cache[key] = cached - self._construction_dirty = False - except BaseException as error: - if cached is None or self._plan_cache.get(key) is not cached: - self._reconcile_after_construction_error_unlocked( - error, "indexed plan" - ) - else: - # Exact cache identity proves the STORE committed - # even when its return boundary was interrupted. - self._construction_dirty = False - raise - # Publish the plan and immutable index ownership in one state - # transition. Core snapshots and validates the indices during - # its allocation-minimal direct one-shot submit. - handle = _WorkHandle( - operation=operation, - plan=cached.plan, - local_indices=local_indices, - remote_indices=remote_indices, - notification=wire_notification, - peer_name=remote.peer_name, - cache_key=key, - ) - return self._track_work(handle) - - def initialize_xfer( - self, - operation: str, - source: _RawDescriptorList, - destination: _RawDescriptorList, - peer_name: str, - notification: bytes = b"", - ) -> _WorkHandle: - with self._lock: - self._require_endpoint() - operation = self._normalize_operation(operation) - if peer_name in self._retiring_peers: - raise RuntimeError(f"Peer {peer_name!r} teardown is in progress") - if operation == "WRITE": - local_regions = self._resolve_local( - source.descriptors, source.memory_type - ) - remote_regions = self._resolve_remote( - peer_name, destination.descriptors, destination.memory_type - ) - else: - local_regions = self._resolve_local( - destination.descriptors, destination.memory_type - ) - remote_regions = self._resolve_remote( - peer_name, source.descriptors, source.memory_type - ) - self._reconcile_construction_state_unlocked() - self._construction_dirty = True - plan = None - handle = None - try: - plan = self._require_endpoint().prepare( - local=local_regions, - remote=remote_regions, - indexed=False, - ) - handle = _WorkHandle( - owned_plan=plan, - operation=operation, - plan=plan, - notification=bytes(notification) if notification else None, - peer_name=peer_name, - ) - tracked = self._track_work(handle) - self._construction_dirty = False - return tracked - except BaseException as error: - if handle is not None and self._works.get(id(handle)) is handle: - try: - self._release_work(handle) - except BaseException as cleanup_error: - self._retained_work_ids.add(id(handle)) - self._start_retained_work_sweeper() - self._add_exception_note( - error, - "one-off handle cleanup also raised: " - f"{type(cleanup_error).__qualname__}", - ) - self._reconcile_after_construction_error_unlocked(error, "one-off plan") - raise - - def transfer(self, handle: _WorkHandle) -> str: - with self._admitted_call() as endpoint: - should_post = False - with handle._state_lock: - if handle.terminal_state is not None: - return handle.terminal_state - if handle.phase == "released": - raise RuntimeError("PyTorch transfer handle is released") - if handle.phase == "posting": - raise RuntimeError( - "PyTorch transfer submission is already in progress" - ) - if handle.phase == "prepared": - # Publish submission ownership before native entry. The lock - # is released before submit_and_poll/write/read. - handle.phase = "posting" - should_post = True - recovery_pending = handle.phase == "recovery_pending" - if recovery_pending: - self._retry_interrupted_post_recovery(handle, schedule=True) - with handle._state_lock: - recovery_pending = handle.phase == "recovery_pending" - if recovery_pending: - return "PROC" - if should_post: - try: - self._post(handle, endpoint) - except BaseException as error: - # Recover provider ownership before any diagnostic formatting: - # an exception may itself have hostile __str__/repr hooks. - with handle._state_lock: - native_submit_entered = handle._native_submit_entered - if native_submit_entered: - self._handle_interrupted_submission(handle, error) - else: - # Validation/teardown rejection precedes provider entry, - # so there is no hidden Work to reconcile. - with handle._state_lock: - handle.terminal_state = "ERR" - handle.phase = "terminal" - self._release_work(handle) - with self._lock: - retained = id(handle) in self._works - with handle._state_lock: - retained_phase = handle.phase - if retained and retained_phase in ( - "posted", - "recovery_pending", - ): - # Preserve immutable diagnostics only. Retaining the - # exception would retain this stack through a long DMA. - try: - details = (type(error).__qualname__, str(error)) - except BaseException: - details = ( - "BaseException", - "diagnostic formatting unavailable", - ) - try: - handle.__dict__["_submission_error"] = details - except BaseException: # noqa: S110 - pass - # Return the already tracked handle to conn.py's source - # lifetime barrier. The error is surfaced after terminality. - return "PROC" - raise - state = self._work_state(handle, release=False) - if state == "ERR": - # Current conn.py callers raise immediately on a post failure - # and never issue the later release call. Terminal failure is - # already authoritative, so release it here; a failed cleanup - # remains tracked and is retried by the retained-work sweeper. - try: - self._release_work(handle) - except Exception as exc: - logger.warning("PyTorch failed Work cleanup failed: %s", exc) - with self._lock: - self._retained_work_ids.add(id(handle)) - self._start_retained_work_sweeper() - return state - - def check_xfer_state(self, handle: _WorkHandle) -> str: - with self._admitted_call(): - return self._work_state(handle, release=True) - - def begin_handle_batch(self, handles: list[_WorkHandle]) -> None: - """Publish a worker-owned sink before that worker can create a Work.""" - with self._lock: - self._require_endpoint() - if getattr(self._batch_handles, "value", None) is not None: - raise RuntimeError("PyTorch transfer handle batch is already active") - batch_id = id(handles) - existing = self._active_handle_batches.get(batch_id) - if existing is not None and existing is not handles: - raise RuntimeError("PyTorch transfer handle batch identity collision") - # Publish the list itself before any Core call. Every later handle - # append is therefore already reachable if its worker is interrupted. - try: - self._active_handle_batches[batch_id] = handles - self._batch_handles.value = handles - except BaseException: - # conn.py has not observed a successful begin and therefore - # cannot own cleanup yet. Roll back either store even when the - # store itself committed before raising. - if getattr(self._batch_handles, "value", None) is handles: - del self._batch_handles.value - if self._active_handle_batches.get(batch_id) is handles: - self._active_handle_batches.pop(batch_id, None) - raise - - def end_handle_batch(self, handles: list[_WorkHandle]) -> None: - with self._lock: - self._require_endpoint() - current = getattr(self._batch_handles, "value", None) - if current is not handles: - raise RuntimeError("PyTorch transfer handle batch ownership mismatch") - del self._batch_handles.value - self._active_handle_batches.pop(id(handles), None) - self._prune_retained_handle_batches_unlocked() - if id(handles) in self._retained_handle_batches: - self._start_retained_work_sweeper() - - def retain_handle_batch(self, handles: list[_WorkHandle]) -> None: - """Durably retain a whole worker batch with one identity publication.""" - - with self._lock: - self._require_endpoint() - active = self._active_handle_batches.get(id(handles)) - if active is not None and active is not handles: - raise RuntimeError("PyTorch transfer handle batch identity collision") - self._retained_handle_batches[id(handles)] = handles - self._prune_retained_handle_batches_unlocked() - if id(handles) in self._retained_handle_batches: - self._start_retained_work_sweeper() - - def handle_batch_is_quiescent(self, handles: Sequence[_WorkHandle]) -> bool: - with self._lock: - return not any(self._works.get(id(handle)) is handle for handle in handles) - - @staticmethod - def pop_xfer_error(handle: _WorkHandle) -> RuntimeError | None: - """Return deferred lost-return diagnostics after terminal settlement.""" - - details = handle.__dict__.pop("_submission_error", None) - if details is None: - return None - error_type, message = details - return RuntimeError( - "PyTorch transfer submission return was interrupted by " - f"{error_type}: {message}" - ) - - def progress(self) -> None: - """Advance the provider and reap detached notification requests. - - SGLang calls this before checking one batch of transfer handles. Regular - Works are intentionally not scanned here: the following per-handle loop - is their single status pass. Detached notification Works have no caller - to check them, so this method still owns their terminal cleanup. - """ - - # Optimistic empty fast path: a concurrent sender publishes the Work in - # ``_works`` before appending here, so missing that append only defers - # detached cleanup to the next poll/close and cannot lose ownership. - if self._background_progress and not self._detached_works: - return - with self._admitted_call(): - self._advance_provider() - self._reap_detached_works() - - def cancel_handles(self, handles: Sequence[_WorkHandle]) -> None: - """Best-effort cancel siblings while retaining non-terminal Works.""" - - with self._admitted_call(): - for handle in handles: - with self._lock: - tracked = id(handle) in self._works - if not tracked: - continue - with handle._state_lock: - phase = handle.phase - if phase == "prepared": - handle.terminal_state = "ERR" - handle.phase = "terminal" - if phase == "prepared": - try: - self._release_work(handle) - except Exception as exc: - logger.warning( - "PyTorch prepared transfer cleanup failed: %s", exc - ) - with self._lock: - self._retained_work_ids.add(id(handle)) - self._start_retained_work_sweeper() - continue - if phase == "posting": - with self._lock: - self._retained_work_ids.add(id(handle)) - continue - if phase == "recovery_pending": - self._retry_interrupted_post_recovery(handle, schedule=True) - with handle._state_lock: - recovery_pending = handle.phase == "recovery_pending" - if recovery_pending: - continue - if self._work_state(handle, release=False) in ("DONE", "ERR"): - continue - try: - with handle._state_lock: - work = handle.work - assert work is not None - work.cancel() - except Exception as exc: - logger.warning("PyTorch transfer cancellation failed: %s", exc) - self._progress() - with self._lock: - tracked_handles = tuple( - handle - for handle in handles - if self._works.get(id(handle)) is handle - ) - retained = set() - for handle in tracked_handles: - with handle._state_lock: - recovery_pending = handle.phase == "recovery_pending" - if recovery_pending or self._work_state(handle, release=False) not in ( - "DONE", - "ERR", - ): - retained.add(id(handle)) - if retained: - with self._lock: - self._retained_work_ids.update(retained) - self._start_retained_work_sweeper() - - def cancel_handles_and_wait(self, handles: Sequence[_WorkHandle]) -> None: - """Cancel a failed batch and return only when every Work is reusable. - - This deliberately unbounded routine is used only on failure. Publishing - a room failure while an uncancellable Work is active would let SGLang - recycle KV pages that DMA may still be reading. - """ - - if isinstance(handles, list): - self.retain_handle_batch(handles) - handles = tuple(handles) - with self._admitted_call(): - self.cancel_handles(handles) - while True: - pending = False - self._advance_provider() - for handle in handles: - with self._lock: - tracked = self._works.get(id(handle)) is handle - if not tracked: - continue - with handle._state_lock: - phase = handle.phase - if phase == "recovery_pending": - self._retry_interrupted_post_recovery( - handle, - schedule=False, - ) - with handle._state_lock: - phase = handle.phase - if phase in ("posting", "recovery_pending"): - pending = True - continue - if phase == "prepared": - with handle._state_lock: - if handle.phase == "prepared": - handle.terminal_state = "ERR" - handle.phase = "terminal" - if self._work_state(handle, release=False) not in ( - "DONE", - "ERR", - ): - pending = True - continue - try: - self._release_work(handle) - except Exception as exc: - logger.warning( - "PyTorch failed-batch terminal cleanup failed: %s", - exc, - ) - pending = True - if not pending: - with self._lock: - self._prune_retained_handle_batches_unlocked() - return - time.sleep(0) - - def release_xfer_handle(self, handle: _WorkHandle) -> None: - with self._admitted_call(): - with self._lock: - tracked = self._works.get(id(handle)) is handle - if not tracked: - return - with handle._state_lock: - phase = handle.phase - if phase == "prepared": - self._release_work(handle) - return - if phase == "posting": - raise RuntimeError("Cannot release a posting PyTorch transfer") - if phase == "recovery_pending": - self._retry_interrupted_post_recovery(handle, schedule=True) - with handle._state_lock: - recovery_pending = handle.phase == "recovery_pending" - if recovery_pending: - raise RuntimeError( - "Cannot release a PyTorch transfer while Work recovery is pending" - ) - if self._work_state(handle, release=False) not in ("DONE", "ERR"): - with handle._state_lock: - work = handle.work - assert work is not None - work.cancel() - self._progress() - if self._work_state(handle, release=False) not in ("DONE", "ERR"): - raise RuntimeError("Cannot release an active PyTorch transfer") - self._release_work(handle) - - def send_notif(self, peer_name: str, payload: bytes) -> _WorkHandle | None: - with self._admitted_call() as endpoint: - with self._lock: - peer = self._require_peer(peer_name) - send_notification_sync = self._send_notification_sync - if type(payload) is not bytes: - payload = bytes(payload) - if send_notification_sync is not None: - send_notification_sync(peer, payload) - return None - handle = _WorkHandle(phase="posting") - try: - # Tracking and provider entry are one recovery transaction. - self._track_work(handle) - work = endpoint.send_notification( - peer, - payload, - recovery_token=handle, - ) - with handle._state_lock: - handle.work = work - handle.phase = "posted" - with self._lock: - self._detached_works.append(handle) - return handle - except BaseException as error: - self._handle_interrupted_submission(handle, error) - with self._lock: - tracked = self._works.get(id(handle)) is handle - if tracked: - self._ensure_detached(handle) - raise - - def _notification_alias( - self, - source_endpoint_id: str | None, - source_incarnation: str | None, - source_aliases: Mapping[object, str] | None = None, - ) -> str: - if source_endpoint_id is None: - if source_incarnation is not None: - raise RuntimeError("notification source identity is incomplete") - # One-way WRITE topologies need not import the sender. Notification - # payloads remain self-routing in SGLang's current protocol. - return _UNKNOWN_NOTIFICATION_SOURCE - if source_incarnation is None: - raise RuntimeError( - "notification source is missing its endpoint incarnation" - ) - source_key = (str(source_endpoint_id), str(source_incarnation)) - if source_aliases is None: - with self._lock: - source_aliases = self._source_aliases - try: - return source_aliases[source_key] - except KeyError as exc: - raise RuntimeError( - "notification arrived from an unknown or stale endpoint " - f"incarnation {source_key!r}" - ) from exc - - def get_new_notifs(self) -> dict[str, list[bytes]]: - with self._admitted_call() as endpoint: - poll_batches = getattr(endpoint, "poll_notification_batches", None) - if callable(poll_batches): - notifications = poll_batches() - grouped_receive = True - else: - notifications = endpoint.poll_notifications() - grouped_receive = False - # Either polling operation already drives NIXL notification progress. - # Reap detached notification Works without issuing a second - # provider progress call or scanning caller-owned transfers. - self._reap_detached_works() - # Peer mutation publishes a replacement alias map under ``_lock``. - # Snapshot only after native polling, then route the whole batch - # without a lock per notification. - with self._lock: - source_aliases = self._source_aliases - grouped: dict[str, list[bytes]] = defaultdict(list) - if grouped_receive: - for batch in notifications: - alias = self._notification_alias( - batch.source_endpoint_id, - batch.source_incarnation, - source_aliases, - ) - grouped[alias].extend(batch.payloads) - return dict(grouped) - if isinstance(notifications, Mapping): - for source, payloads in notifications.items(): - if source is None: - alias = _UNKNOWN_NOTIFICATION_SOURCE - else: - source = str(source) - alias = source_aliases.get(source, source) - grouped[alias].extend(bytes(payload) for payload in payloads) - return dict(grouped) - for notification in notifications: - alias = self._notification_alias( - notification.source_endpoint_id, - getattr(notification, "source_incarnation", None), - source_aliases, - ) - grouped[alias].append(bytes(notification.payload)) - return dict(grouped) - - def close(self) -> None: - endpoint = self._begin_close() - if endpoint is None: - return - try: - # Endpoint.closed is the authoritative commit receipt. Inspect it - # before querying children because a previous close may have - # committed even though its Python return was lost. - if self._endpoint_is_closed(endpoint): - self._finalize_closed_endpoint_unlocked(endpoint) - return - self._close_open_endpoint_exclusive(endpoint) - finally: - self._end_close() - - def _close_open_endpoint_exclusive(self, endpoint: object) -> None: - """Close one live endpoint while the lifecycle gate is exclusive.""" - - try: - self._retry_endpoint_cleanup_unlocked() - self._reconcile_construction_state_unlocked(force=True) - for handle in tuple(self._works.values()): - if handle.phase == "recovery_pending": - self._retry_interrupted_post_recovery(handle, schedule=False) - for handle in tuple(self._works.values()): - if handle.phase == "prepared": - self._release_work(handle) - for handle in tuple(self._works.values()): - if handle.phase == "posted" and self._work_state( - handle, release=False - ) not in ("DONE", "ERR"): - assert handle.work is not None - handle.work.cancel() - self._progress() - active = [ - handle - for handle in self._works.values() - if handle.phase in ("posting", "recovery_pending") - or ( - handle.phase == "posted" - and self._work_state(handle, release=False) not in ("DONE", "ERR") - ) - ] - if active: - raise RuntimeError( - f"Cannot close endpoint with {len(active)} active transfer(s)" - ) - for handle in tuple(self._works.values()): - self._release_work(handle) - self._detached_works.clear() - self._retained_work_ids.clear() - self._active_handle_batches.clear() - self._retained_handle_batches.clear() - self._retiring_peers.update(self._peers) - self._release_cached_plans(lambda cached: True) - for peer in reversed(tuple(self._peers.values())): - peer.close() - self._peers.clear() - self._peer_wire_metadata.clear() - self._retiring_peers.clear() - for record in reversed(self._registrations): - record.registration.close() - self._registrations.clear() - self._registration_groups.clear() - endpoint.close() - except BaseException: - if self._endpoint_is_closed(endpoint): - self._finalize_closed_endpoint_unlocked(endpoint) - raise - self._finalize_closed_endpoint_unlocked(endpoint) - - @contextmanager - def _admitted_call(self): - """Admit native work without holding a Python lock across the call.""" - - depth = getattr(self._admission_local, "depth", 0) - if depth: - self._admission_local.depth = depth + 1 - try: - yield self._require_endpoint() - finally: - self._admission_local.depth = depth - return - thread_id = threading.get_ident() - with self._lifecycle_cv: - if self._closing and self._close_owner != thread_id: - raise RuntimeError("PyTorch transfer endpoint close is in progress") - endpoint = self._require_endpoint() - self._active_calls += 1 - self._admission_local.depth = 1 - try: - yield endpoint - finally: - with self._lifecycle_cv: - self._admission_local.depth = 0 - self._active_calls -= 1 - if self._active_calls == 0: - self._lifecycle_cv.notify_all() - - def _begin_close(self) -> object | None: - thread_id = threading.get_ident() - with self._lifecycle_cv: - while self._closing and self._close_owner != thread_id: - self._lifecycle_cv.wait() - if self._close_owner == thread_id: - raise RuntimeError("recursive PyTorch transfer endpoint close") - if self._endpoint is None: - return None - self._closing = True - self._close_owner = thread_id - while self._active_calls: - self._lifecycle_cv.wait() - return self._endpoint - - def _end_close(self) -> None: - with self._lifecycle_cv: - self._closing = False - self._close_owner = None - self._lifecycle_cv.notify_all() - - @staticmethod - def _endpoint_is_closed(endpoint: object) -> bool: - try: - return getattr(endpoint, "closed", False) is True - except BaseException: - return False - - def _finalize_closed_endpoint_unlocked(self, endpoint: object) -> None: - """Adopt authoritative endpoint close without touching Core children.""" - - if self._endpoint is not endpoint: - return - for handle in self._works.values(): - if handle.terminal_state is None: - handle.terminal_state = "ERR" - handle.phase = "released" - handle.local_indices = None - handle.remote_indices = None - handle.notification = None - handle.plan = None - self._works.clear() - self._detached_works.clear() - self._retained_work_ids.clear() - self._active_handle_batches.clear() - self._retained_handle_batches.clear() - self._plan_cache.clear() - self._closing_plans.clear() - self._retiring_peers.clear() - self._peers.clear() - self._peer_wire_metadata.clear() - self._source_aliases.clear() - self._registrations.clear() - self._registration_groups.clear() - self._send_notification_sync = None - self._endpoint = None - self._construction_dirty = False - self._sweep_wakeup.set() - - @staticmethod - def _supports_capability(endpoint: object, name: str) -> bool: - capabilities = getattr(endpoint, "capabilities", None) - if callable(capabilities): - capabilities = capabilities() - if isinstance(capabilities, Mapping): - return bool(capabilities.get(name, False)) - return bool(getattr(capabilities, name, False)) - - @staticmethod - def _normalize_descriptors( - descriptors: Sequence[Sequence[object]], - ) -> tuple[_RawDescriptor, ...]: - normalized = [] - for descriptor in descriptors: - if len(descriptor) < 3: - raise ValueError("Transfer descriptors require address, size, device") - address, nbytes, device_id = map(int, descriptor[:3]) - if address <= 0 or nbytes <= 0 or device_id < 0: - raise ValueError( - "Transfer descriptor address/size must be positive and device " - "must be non-negative" - ) - stride: int | None = None - count: int | None = None - if len(descriptor) >= 5: - stride, count = map(int, descriptor[3:5]) - if stride < nbytes or count <= 0: - raise ValueError( - "Strided descriptors require stride >= size and count > 0" - ) - if count == 1: - stride = count = None - elif len(descriptor) == 4 and not isinstance(descriptor[3], str): - raise ValueError("Strided descriptors require both stride and count") - normalized.append(_RawDescriptor(address, nbytes, device_id, stride, count)) - return tuple(normalized) - - @staticmethod - def _device(memory_type: str, device_id: int) -> str: - if memory_type == "DRAM": - return "cpu" - if memory_type == "VRAM": - return f"cuda:{device_id}" - raise NotImplementedError( - f"PyTorch transfer adapter does not support {memory_type!r} memory" - ) - - def _resolve_local( - self, descriptors: Sequence[_RawDescriptor], memory_type: str - ) -> tuple[object, ...]: - regions = [] - for descriptor in descriptors: - end = descriptor.address + descriptor.extent - for record in self._registrations: - if ( - record.memory_type == memory_type - and record.device_id == descriptor.device_id - and descriptor.address >= record.address - and end <= record.address + record.nbytes - ): - kwargs = ( - {} - if descriptor.count is None - else { - "stride": descriptor.stride, - "count": descriptor.count, - } - ) - regions.append( - record.registration.region( - descriptor.address - record.address, - descriptor.nbytes, - **kwargs, - ) - ) - break - else: - raise ValueError( - f"Local raw span 0x{descriptor.address:x}+{descriptor.nbytes} " - "is outside registered memory" - ) - return tuple(regions) - - def _resolve_remote( - self, - peer_name: str, - descriptors: Sequence[_RawDescriptor], - memory_type: str, - ) -> tuple[object, ...]: - peer = self._require_peer(peer_name) - return tuple( - peer.unsafe_region( - descriptor.address, - descriptor.nbytes, - device=self._device(memory_type, descriptor.device_id), - memory_type=memory_type, - **( - {} - if descriptor.count is None - else { - "stride": descriptor.stride, - "count": descriptor.count, - } - ), - ) - for descriptor in descriptors - ) - - def _orient_prepared( - self, - operation: str, - source: _PreparedList, - source_indices: Sequence[int], - destination: _PreparedList, - destination_indices: Sequence[int], - ) -> tuple[_PreparedList, _PreparedList, Sequence[int], Sequence[int]]: - if operation == "WRITE": - local, remote = source, destination - local_indices, remote_indices = source_indices, destination_indices - else: - local, remote = destination, source - local_indices, remote_indices = destination_indices, source_indices - if local.remote or not remote.remote: - raise ValueError( - f"Invalid {operation} prepared-list local/remote orientation" - ) - return local, remote, local_indices, remote_indices - - def _post(self, handle: _WorkHandle, endpoint: object) -> None: - if handle.operation is None or handle.plan is None: - raise RuntimeError("PyTorch transfer handle has no prepared operation") - with self._lock: - if handle.peer_name in self._retiring_peers: - raise RuntimeError(f"Peer {handle.peer_name!r} teardown is in progress") - if handle.peer_name is not None: - self._require_peer(handle.peer_name) - if handle.cache_key is not None and handle.cache_key in self._closing_plans: - raise RuntimeError("Prepared transfer plan teardown is in progress") - with handle._state_lock: - handle._native_submit_entered = True - first_state = None - if handle.local_indices is not None or handle.remote_indices is not None: - if handle.local_indices is None or handle.remote_indices is None: - raise RuntimeError("PyTorch indexed transfer has incomplete indices") - kwargs = { - "local_indices": handle.local_indices, - "remote_indices": handle.remote_indices, - "max_polls": 0, - "timeout_ns": None, - "recovery_token": handle, - } - if handle.notification is None: - work, first_state = handle.plan.submit_indices_and_poll( - self._core_transfer_operations[handle.operation], - **kwargs, - ) - else: - work, first_state = handle.plan.submit_indices_and_poll( - self._core_transfer_operations[handle.operation], - **kwargs, - notification=handle.notification, - ) - else: - method = getattr(endpoint, handle.operation.lower()) - if handle.notification is None: - work = method( - handle.plan, - recovery_token=handle, - ) - else: - work = method( - handle.plan, - notification=handle.notification, - recovery_token=handle, - ) - with handle._state_lock: - handle.work = work - handle.first_state = first_state - # Core Work now owns its authenticated private snapshot. Drop the - # adapter's temporary immutable inputs on the successful boundary. - handle.local_indices = None - handle.remote_indices = None - handle.phase = "posted" - - def _recover_interrupted_post_work( - self, - handle: _WorkHandle, - submission_error: BaseException | None = None, - ) -> Literal["absent", "recovered", "pending"]: - """Recover one Core Work by an exact, caller-owned identity token.""" - - # Claim recovery with a state-only critical section. The Core ownership - # snapshot below must not run under either an agent or handle lock, and - # a second recovery caller must not race adoption of the same Work. - with handle._state_lock: - if handle.work is not None: - handle.phase = "posted" - return "recovered" - if handle.phase not in ("posting", "recovery_pending"): - return "absent" - if handle._recovery_in_progress: - return "pending" - handle.phase = "recovery_pending" - handle._recovery_in_progress = True - match = None - matches = () - try: - # This snapshot stays entirely off the successful path. Distractor - # Works cannot be adopted: token comparison is strictly ``is``. - matches = tuple( - work - for work in self._require_endpoint().open_works - if work.recovery_token is handle - ) - except BaseException as recovery_error: - self._note_recovery_issue(submission_error, recovery_error) - recovery = "pending" - else: - if not matches: - recovery = "absent" - elif len(matches) == 1: - match = matches[0] - recovery = "recovered" - else: - recovery = "pending" - if len(matches) > 1: - self._note_recovery_issue( - submission_error, - RuntimeError( - "multiple Core Works carry one SGLang recovery-token identity" - ), - ) - with handle._state_lock: - try: - if recovery == "recovered": - handle.work = match - handle.phase = "posted" - elif recovery == "absent": - # Absence is authoritative only for the recovery owner. Mark - # terminal before releasing the claim so no retry can race - # between the empty snapshot and terminal publication. - handle.terminal_state = "ERR" - handle.phase = "terminal" - finally: - handle._recovery_in_progress = False - return recovery - - def _handle_interrupted_submission( - self, - handle: _WorkHandle, - submission_error: BaseException, - ) -> None: - with handle._state_lock: - handle.submission_interrupted = ( - handle.phase - in ( - "posting", - "recovery_pending", - "posted", - ) - or handle.work is not None - ) - recovery = self._recover_interrupted_post_work(handle, submission_error) - if recovery == "absent": - self._close_failed_submission( - handle, - submission_error, - schedule=True, - ) - return - if recovery == "recovered": - self._quarantine_recovered_submission(handle, schedule=True) - return - self._retain_recovery_pending(handle, schedule=True) - - def _retry_interrupted_post_recovery( - self, - handle: _WorkHandle, - *, - schedule: bool, - ) -> None: - with handle._state_lock: - if handle.phase != "recovery_pending": - return - recovery = self._recover_interrupted_post_work(handle) - if recovery == "pending": - self._retain_recovery_pending(handle, schedule=schedule) - return - if recovery == "absent": - self._close_failed_submission(handle, schedule=schedule) - return - self._quarantine_recovered_submission(handle, schedule=schedule) - - def _quarantine_recovered_submission( - self, - handle: _WorkHandle, - *, - schedule: bool, - ) -> None: - with handle._state_lock: - work = handle.work - assert work is not None - try: - work.cancel() - except Exception as exc: - logger.warning("PyTorch interrupted Work cancellation failed: %s", exc) - if self._work_state(handle, release=False) in ("DONE", "ERR"): - self._close_failed_submission(handle, schedule=schedule) - return - self._retain_recovery_pending(handle, schedule=schedule) - - def _close_failed_submission( - self, - handle: _WorkHandle, - submission_error: BaseException | None = None, - *, - schedule: bool, - ) -> None: - try: - self._release_work(handle) - except BaseException as cleanup_error: - self._note_recovery_issue(submission_error, cleanup_error) - logger.warning("PyTorch failed-post cleanup failed: %s", cleanup_error) - self._retain_recovery_pending(handle, schedule=schedule) - if not isinstance(cleanup_error, Exception): - raise - - def _retain_recovery_pending( - self, - handle: _WorkHandle, - *, - schedule: bool, - ) -> None: - retained = False - with self._lock: - if self._works.get(id(handle)) is handle: - self._retained_work_ids.add(id(handle)) - retained = True - if schedule and retained: - self._start_retained_work_sweeper() - - def _require_no_unsettled_interrupted_submission(self, action: str) -> None: - """Fail closed using state only; callers may already own ``_lock``. - - Recovery and status observation are driven by admitted transfer or - sweeper calls. A teardown path must never hide either Core operation - beneath the agent RLock merely to make cleanup opportunistically pass. - """ - - unsettled = [] - for handle in tuple(self._works.values()): - with handle._state_lock: - phase = handle.phase - submission_interrupted = handle.submission_interrupted - if phase in ("posting", "recovery_pending"): - unsettled.append(handle) - continue - if not submission_interrupted or phase not in ( - "posted", - "terminal", - ): - continue - unsettled.append(handle) - if unsettled: - raise RuntimeError( - f"Cannot {action} with {len(unsettled)} interrupted PyTorch " - "submission(s) unsettled" - ) - - def _ensure_detached(self, handle: _WorkHandle) -> None: - with self._lock: - if self._works.get(id(handle)) is not handle: - return - if not any(candidate is handle for candidate in self._detached_works): - self._detached_works.append(handle) - - @staticmethod - def _note_recovery_issue( - submission_error: BaseException | None, - recovery_error: BaseException, - ) -> None: - if submission_error is None: - return - add_note = getattr(submission_error, "add_note", None) - if callable(add_note): - try: - add_note( - "PyTorch interrupted-post Work recovery reported " - f"{type(recovery_error).__qualname__}" - ) - except BaseException: # noqa: S110 - pass - - def _track_work(self, handle: _WorkHandle) -> _WorkHandle: - batch_handles = getattr(self._batch_handles, "value", None) - try: - with self._lock: - if batch_handles is not None and handle.operation is not None: - # This publication precedes provider entry. If a later helper - # frame is interrupted, the room-level barrier still owns it. - batch_handles.append(handle) - # Publish to the graph only after a worker batch can reach it. - self._works[id(handle)] = handle - return handle - except BaseException: - with self._lock: - self._works.pop(id(handle), None) - if batch_handles is not None: - batch_handles[:] = [ - candidate - for candidate in batch_handles - if candidate is not handle - ] - raise - - def _progress(self) -> None: - self._advance_provider() - self._reap_works() - - def _advance_provider(self) -> None: - endpoint = self._require_endpoint() - if not self._background_progress: - try: - endpoint.progress() - except Exception as exc: - # Global progress is an observation boundary too: failure does - # not prove that any request is terminal. Subsequent per-Work - # probes retain every unknown request and recover independently. - logger.warning( - "PyTorch transfer progress is temporarily unknown: %s", exc - ) - - def _reap_works(self) -> None: - with self._lock: - handles = tuple(self._works.values()) - self._reap_handles(handles) - - def _reap_detached_works(self) -> None: - with self._lock: - handles = tuple(self._detached_works) - if not handles: - return - self._reap_handles(handles) - - def _reap_retained_works(self) -> None: - retained = [] - seen = set() - with self._lock: - for handle_id in self._retained_work_ids: - handle = self._works.get(handle_id) - if handle is not None: - retained.append(handle) - seen.add(handle_id) - for handles in self._retained_handle_batches.values(): - for handle in handles: - handle_id = id(handle) - if handle_id in seen or self._works.get(handle_id) is not handle: - continue - retained.append(handle) - seen.add(handle_id) - for handle in retained: - with handle._state_lock: - if handle.phase == "prepared": - handle.terminal_state = "ERR" - handle.phase = "terminal" - self._reap_handles(tuple(retained)) - with self._lock: - self._prune_retained_handle_batches_unlocked() - - def _reap_handles(self, handles: Sequence[_WorkHandle]) -> None: - for handle in handles: - with self._lock: - tracked = self._works.get(id(handle)) is handle - if not tracked: - continue - with handle._state_lock: - phase = handle.phase - if phase == "recovery_pending": - self._retry_interrupted_post_recovery(handle, schedule=False) - with handle._state_lock: - phase = handle.phase - if phase in ("prepared", "posting", "recovery_pending"): - continue - if self._work_state(handle, release=False) not in ("DONE", "ERR"): - continue - try: - self._release_work(handle) - except Exception as exc: - # The Work remains tracked, and the next progress call retries - # cleanup without losing its terminal state. - logger.warning("PyTorch terminal Work cleanup failed: %s", exc) - with self._lock: - self._detached_works = [ - handle - for handle in self._detached_works - if self._works.get(id(handle)) is handle - ] - self._retained_work_ids.intersection_update(self._works) - - def _start_retained_work_sweeper(self) -> None: - with self._lock: - if self._sweep_thread is not None and self._sweep_thread.is_alive(): - self._sweep_wakeup.set() - return - self._sweep_wakeup.clear() - self._sweep_thread = threading.Thread( - target=self._sweep_retained_works, - name=f"{self.name}-retained-work-sweeper", - daemon=True, - ) - self._sweep_thread.start() - - def _sweep_retained_works(self) -> None: - """Periodically reap abandoned siblings without freeing active Works.""" - - while True: - self._sweep_wakeup.wait(_RETAINED_WORK_SWEEP_INTERVAL_S) - self._sweep_wakeup.clear() - with self._lock: - if self._endpoint is None: - self._sweep_thread = None - return - self._retained_work_ids.intersection_update(self._works) - self._prune_retained_handle_batches_unlocked() - if not self._retained_work_ids and not self._retained_handle_batches: - self._sweep_thread = None - return - try: - # Abandoned siblings need autonomous provider progress, but no - # adapter lock spans provider progress or Work-state refresh. - with self._admitted_call(): - self._advance_provider() - self._reap_retained_works() - self._reap_detached_works() - except RuntimeError: - with self._lock: - if self._closing or self._endpoint is None: - self._sweep_thread = None - return - raise - - def _prune_retained_handle_batches_unlocked(self) -> None: - for batch_id, handles in tuple(self._retained_handle_batches.items()): - if not any(self._works.get(id(handle)) is handle for handle in handles): - self._retained_handle_batches.pop(batch_id, None) - - def _work_state(self, handle: _WorkHandle, *, release: bool) -> str: - while True: - with handle._state_lock: - terminal_state = handle.terminal_state - phase = handle.phase - if terminal_state is not None: - break - if phase == "prepared": - raise RuntimeError("PyTorch transfer has not been posted") - if phase == "posting": - raise RuntimeError( - "PyTorch transfer submission is still in progress" - ) - if phase == "recovery_pending": - recover = True - else: - recover = False - if handle.work is None: - raise RuntimeError("Posted PyTorch transfer has no Work") - if handle._status_in_progress: - return "PROC" - handle._status_in_progress = True - work = handle.work - state = handle.first_state - handle.first_state = None - if not recover: - break - self._retry_interrupted_post_recovery(handle, schedule=True) - with handle._state_lock: - recovery_pending = handle.phase == "recovery_pending" - if recovery_pending: - return "PROC" - - if terminal_state is not None: - if release: - self._release_work(handle) - return terminal_state - try: - # Work.state is the sole authoritative observation. In current - # Core it refreshes once and already incorporates provider errors; - # reading test(), error, and state separately would refresh the - # same Work up to three times. - if state is None: - state = work.state - outcomes = self._work_state_outcomes - if outcomes is not None: - try: - terminal = outcomes.get(state, _UNKNOWN_WORK_OUTCOME) - except TypeError: - terminal = _UNKNOWN_WORK_OUTCOME - if terminal is _UNKNOWN_WORK_OUTCOME: - raise ValueError(f"Unknown PyTorch WorkState: {state!r}") - if terminal is None: - return "PROC" - else: - state_name = self._state_name(state) - if state_name in _PENDING_STATES: - return "PROC" - if state_name in _DONE_STATES: - terminal = "DONE" - elif state_name in _FAILED_STATES: - terminal = "ERR" - else: - raise ValueError(f"Unknown PyTorch WorkState: {state!r}") - except Exception as exc: - # A status observation failure does not prove the transport is - # terminal. Returning PROC preserves all memory/plan lifetimes and - # permits a later poll to recover or report a real terminal error. - logger.warning("PyTorch transfer status is temporarily unknown: %s", exc) - return "PROC" - finally: - with handle._state_lock: - handle._status_in_progress = False - with handle._state_lock: - if handle.submission_interrupted: - terminal = "ERR" - if handle.terminal_state is None: - handle.terminal_state = terminal - handle.phase = "terminal" - else: - terminal = handle.terminal_state - if release: - self._release_work(handle) - return terminal - - @staticmethod - def _complete_work_state_outcomes( - api: Any, - ) -> dict[object, str | None] | None: - work_state = getattr(api, "WorkState", None) - try: - outcomes = { - work_state.PENDING: None, - work_state.RUNNING: None, - work_state.COMPLETED: "DONE", - work_state.FAILED: "ERR", - work_state.CANCELLED: "ERR", - } - except (AttributeError, TypeError): - return None - return outcomes if len(outcomes) == 5 else None - - def _release_work(self, handle: _WorkHandle) -> None: - with self._lock: - tracked = self._works.get(id(handle)) is handle - if not tracked: - return - with handle._state_lock: - if handle.phase == "released" or handle._release_in_progress: - return - if handle.phase in ("posting", "recovery_pending"): - raise RuntimeError( - "Cannot release a Work with unresolved submission ownership" - ) - handle._release_in_progress = True - work = handle.work - owned_plan = handle.owned_plan - try: - if work is not None: - work.close() - if owned_plan is not None: - owned_plan.close() - except BaseException: - with handle._state_lock: - handle._release_in_progress = False - raise - with self._lock: - if self._works.get(id(handle)) is handle: - self._works.pop(id(handle), None) - if handle.operation is None: - self._detached_works = [ - candidate - for candidate in self._detached_works - if candidate is not handle - ] - with handle._state_lock: - handle.phase = "released" - handle.local_indices = None - handle.remote_indices = None - handle.notification = None - handle.plan = None - handle._release_in_progress = False - - def _release_cached_plans(self, predicate: Callable[[_CachedPlan], bool]) -> None: - # Move every matching plan out of the submit-visible cache before the - # first close attempt. A partially committed provider close remains - # discoverable here and may be retried, but can never be submitted. - for key, cached in tuple(self._plan_cache.items()): - if not predicate(cached): - continue - del self._plan_cache[key] - self._closing_plans[key] = cached - for key, cached in reversed(tuple(self._closing_plans.items())): - if not predicate(cached): - continue - if any(handle.cache_key == key for handle in self._works.values()): - raise RuntimeError( - "Prepared transfer plan is still referenced by a transfer handle" - ) - cached.plan.close() - del self._closing_plans[key] - - @staticmethod - def _contains_identity(values: Sequence[object], candidate: object) -> bool: - return any(value is candidate for value in values) - - def _reconcile_construction_state_unlocked(self, *, force: bool = False) -> None: - """Close exact Core children absent from this adapter's owner graph.""" - - if not force and not self._construction_dirty: - return - endpoint = self._require_endpoint() - entries = tuple(self._plan_cache.values()) + tuple(self._closing_plans.values()) - owned_plans = tuple(entry.plan for entry in entries) + tuple( - handle.owned_plan - for handle in self._works.values() - if handle.owned_plan is not None - ) - owned_peers = tuple(self._peers.values()) - owned_registrations = tuple( - record.registration for record in self._registrations - ) - # SGLang does not currently construct bindings or request slots. They - # are still scanned first so any future interrupted child cannot pin a - # lost plan while being silently mistaken for an owned object. - for attribute, owned in ( - ("open_transfer_slots", ()), - ("open_bound_transfers", ()), - ("open_plans", owned_plans), - ("open_peers", owned_peers), - ("open_registrations", owned_registrations), - ): - for child in tuple(getattr(endpoint, attribute)): - if not self._contains_identity(owned, child): - child.close() - self._construction_dirty = False - - def _reconcile_after_construction_error_unlocked( - self, error: BaseException, kind: str - ) -> None: - try: - self._reconcile_construction_state_unlocked(force=True) - except BaseException as cleanup_error: - self._add_exception_note( - error, - f"{kind} construction cleanup also raised: " - f"{type(cleanup_error).__qualname__}", - ) - - def _matches_recovery_endpoint( - self, - candidate: object, - *, - endpoint_id: str, - progress_mode: object, - thread_mode: object, - ) -> bool: - return ( - type(candidate) is self._api.Endpoint - and getattr(candidate, "endpoint_id", None) == endpoint_id - and getattr(candidate, "name", None) == self.name - and getattr(candidate, "backend", None) == self._endpoint_backend - and getattr(candidate, "progress_mode", None) is progress_mode - and getattr(candidate, "thread_mode", None) is thread_mode - ) - - @staticmethod - def _close_recovery_endpoint(endpoint: object) -> None: - if not getattr(endpoint, "closed", False): - endpoint.close() - if not getattr(endpoint, "closed", False): - raise RuntimeError("Endpoint cleanup returned before close committed") - - def _retry_endpoint_cleanup_unlocked(self) -> None: - endpoint = self._endpoint_recovery - if endpoint is not None: - self._close_recovery_endpoint(endpoint) - self._endpoint_recovery = None - - def _retry_endpoint_recovery_unlocked(self) -> None: - """Compatibility alias for focused lifecycle tests.""" - - self._retry_endpoint_cleanup_unlocked() - - @staticmethod - def _add_exception_note(error: BaseException, message: str) -> None: - add_note = getattr(error, "add_note", None) - if callable(add_note): - try: - add_note(message) - except BaseException: # noqa: S110 - pass - - @staticmethod - def _state_name(state: object) -> str | None: - state = getattr(state, "name", state) - return state.upper() if isinstance(state, str) else None - - @staticmethod - def _normalize_operation(operation: str) -> Literal["READ", "WRITE"]: - operation = operation.upper() - if operation not in ("READ", "WRITE"): - raise ValueError(f"Unsupported transfer operation: {operation!r}") - return operation # type: ignore[return-value] - - @staticmethod - def _peer_source_ids(peer: object) -> tuple[str, ...]: - identities = [] - for attribute in ("endpoint_id", "name", "source_endpoint_id"): - value = getattr(peer, attribute, None) - if value is not None: - identity = str(value) - if identity not in identities: - identities.append(identity) - if not identities: - raise TypeError("Imported peer does not expose an endpoint identity") - return tuple(identities) - - def _require_endpoint(self) -> Any: - if self._closing and self._close_owner != threading.get_ident(): - raise RuntimeError("PyTorch transfer endpoint close is in progress") - if self._endpoint is None: - raise RuntimeError("PyTorch transfer endpoint is not initialized") - return self._endpoint - - def _require_peer(self, peer_name: str) -> Any: - self._require_endpoint() - try: - return self._peers[peer_name] - except KeyError as exc: - raise KeyError(f"Unknown PyTorch transfer peer {peer_name!r}") from exc diff --git a/python/sglang/srt/disaggregation/nixl/torch_transfer_mvp.md b/python/sglang/srt/disaggregation/nixl/torch_transfer_mvp.md new file mode 100644 index 000000000000..110498790b7b --- /dev/null +++ b/python/sglang/srt/disaggregation/nixl/torch_transfer_mvp.md @@ -0,0 +1,13 @@ +# Core transfer experiment + +Set `SGLANG_DISAGGREGATION_NIXL_USE_TORCH_TRANSFER=1` on test workers. Requires +the PyTorch `_transfer` prototype and NIXL `torch_transfer` provider. Default: off. + +The NIXL agent is wrapped after creation, retaining its configured backends and +strict thread synchronization. Registration, prepared indexed WRITE, fallback +transfers and notifications route through Core. Bootstrap, allocation ownership, +CUDA ordering, staging and abort protocols are unchanged. + +This is a **success-path review prototype**, not a production mode. The manager +must keep raw-pointer allocations stable. Active cancellation and failure recovery +are outside the MVP. GPU serving/staging and throughput remain unvalidated. diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index a062e1c4b354..d4163ea8835c 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -1191,14 +1191,6 @@ def maybe_send_cached_prefix_chunk(self: Scheduler, req: Req) -> None: if not envs.SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX.get(): return - # Early-send is safe only when the selected data plane can defer source - # access until the transfer worker consumes the recorded CUDA event. - # Keep the capability gate explicit for alternate sender implementations. - if self.enable_overlap and not getattr( - req.disagg_kv_sender, "supports_overlap_early_send", True - ): - return - # Staging sends into positional grid slots, so the early-send boundary # must stay stable across the request's batches: snapshot the at-rest # prefix on the first batch. Non-staging reads the live prefix. diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 272fe38f01a4..3db3f7174058 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -668,19 +668,18 @@ class Envs: # A wedged RDMA stack fails startup here instead of at the scheduler watchdog. SGLANG_DISAGGREGATION_ENGINE_INIT_TIMEOUT = EnvInt(60) SGLANG_DISAGGREGATION_NIXL_BACKEND = EnvStr("UCX") - SGLANG_DISAGGREGATION_NIXL_BACKEND_PARAMS = EnvStr("{}") SGLANG_DISAGGREGATION_NIXL_USE_TORCH_TRANSFER = EnvBool(False) + SGLANG_DISAGGREGATION_NIXL_BACKEND_PARAMS = EnvStr("{}") SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX = EnvBool(True) SGLANG_DISAGGREGATION_ZMQ_MAX_SOCKETS = EnvInt(16384) SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER = EnvBool(False) SGLANG_DISAGGREGATION_FORCE_QUERY_PREFILL_DP_RANK = EnvBool(False) SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS = EnvInt(0) SGLANG_DISAGGREGATION_BOOTSTRAP_ENTRY_CLEANUP_INTERVAL = EnvInt(120) - # Compatibility knob retained for configuration parsing. Decode-side KV - # quarantine is now a mandatory correctness property; a false value is - # ignored by the transfer implementation. The timeout controls only the - # rate-limited operator diagnostic and ABORT retry, never page release. - SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE = EnvBool(True) + # Deferred decode-side KV release: on abort, hold an in-flight request's KV + # pages/slot until the prefill acks the transfer drained, or the timeout + # below fires. Off by default (no behavior/perf impact when disabled). + SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE = EnvBool(False) SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE_TIMEOUT = EnvFloat(30.0) # =================================================================== diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index ad105842f341..764205f08a53 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1068,10 +1068,6 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True): # Cache namespace used to isolate otherwise-identical prefixes. cache_salt: Optional[str] = None - # Internal 128-bit room-lease identity, minted once before scheduler fanout. - # Keep this append-only: TokenizedGenerateReqInput is an array-like wire type. - transfer_generation: Optional[str] = None - def wrap_pickle_fields(self): self.time_stats = wrap_as_pickle(self.time_stats) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index e1173e665fe2..d5e85b986d0f 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -957,7 +957,6 @@ def __init__( disagg_mode: Optional[DisaggregationMode] = None, routed_dp_rank: Optional[int] = None, disagg_prefill_dp_rank: Optional[int] = None, - transfer_generation: Optional[str] = None, vocab_size: Optional[int] = None, priority: Optional[int] = None, metrics_collector: Optional[SchedulerMetricsCollector] = None, @@ -1270,9 +1269,6 @@ def __init__( self.bootstrap_host: str = bootstrap_host self.bootstrap_port: Optional[int] = bootstrap_port self.bootstrap_room: Optional[int] = bootstrap_room - # Shared by every scheduler rank for this decode-room lease. Unlike - # bootstrap_room and rid, this internal value is never user-recycled. - self.transfer_generation: Optional[str] = transfer_generation # Decode-local: the already-emitted boundary token to replay when a # retracted request is rebootstrapped. Set in pause_generation(retract) # and consumed in the decode transfer commit; never plumbed to prefill. diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c6a67140fdb4..85b9c203c950 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2761,7 +2761,6 @@ def handle_generate_request( disagg_mode=self.disaggregation_mode, routed_dp_rank=recv_req.routed_dp_rank, disagg_prefill_dp_rank=recv_req.disagg_prefill_dp_rank, - transfer_generation=recv_req.transfer_generation, vocab_size=self.model_config.vocab_size, priority=recv_req.priority, metrics_collector=( @@ -5240,11 +5239,20 @@ def abort_request(self, recv_req: AbortReq): if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid): logger.debug(f"Abort transfer queue request. {decode_req.req.rid=}") receiver = decode_req.kv_receiver - # Publish exact-generation quarantine before the fallible - # ABORT send. A BaseException/lost return cannot race an ACK - # ahead of decode ownership. - receiver.arm_abort_intent() receiver.abort() + # Arm drain-ack accounting once the ABORT is sent, so acks + # arriving before this req is deferred (e.g. during the next + # forward step) are captured. A fresh set also drops stale acks + # from a prior request that reused this bootstrap_room. A + # redundant abort only re-wipes -- holds longer, never releases + # early -- so no transition guard is needed. + if ( + receiver.kv_mgr.enable_deferred_decode_kv_release + and receiver.abort_notified + ): + receiver.kv_mgr.register_deferred_abort_room( + decode_req.req.bootstrap_room + ) # Abort requests whose KV is already backed up for retraction. if self.disagg_decode_prealloc_queue.retracted_queue: diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 1b3f9a90b88b..ea77294d2f88 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -22,7 +22,6 @@ import logging import os import pickle -import secrets import signal import socket import sys @@ -1427,9 +1426,6 @@ def _create_tokenized_object( return_indexer_topk=obj.return_indexer_topk, routed_dp_rank=obj.routed_dp_rank, disagg_prefill_dp_rank=obj.disagg_prefill_dp_rank, - transfer_generation=( - secrets.token_hex(16) if bootstrap_room is not None else None - ), priority=obj.priority, extra_key=obj.extra_key, cache_salt=obj.cache_salt, diff --git a/rust/sglang-server/src/message/io_struct.rs b/rust/sglang-server/src/message/io_struct.rs index 9cc7fd494a0e..cd8c8fe34842 100644 --- a/rust/sglang-server/src/message/io_struct.rs +++ b/rust/sglang-server/src/message/io_struct.rs @@ -43,9 +43,9 @@ wire_struct! { lora_id: (), custom_logit_processor: (), positional_embed_overrides: (), - /// PD-disaggregation block. The internal generation is append-only at - /// the end of Python's array-like schema, so all intervening defaulted - /// fields must remain present here to preserve their wire positions. + /// PD-disaggregation block — the last fields emitted; everything after + /// `disagg_prefill_dp_rank` in Python has a msgspec default and is + /// omitted (short arrays decode with defaulted tails). bootstrap_host: Option<&'a str>, bootstrap_port: Option, bootstrap_room: Option, @@ -53,21 +53,6 @@ wire_struct! { decode_tp_size: Option, routed_dp_rank: Option, disagg_prefill_dp_rank: Option, - routing_key: (), - require_reasoning: bool, - priority: (), - extra_key: (), - no_logs: bool, - return_bytes: bool, - return_entropy: bool, - need_wait_for_mm_inputs: (), - num_items_assigned: (), - encoder_urls: (), - multi_item_delimiter_indices: (), - time_stats: (), - cache_salt: (), - /// Internal room-lease identity shared by every scheduler rank. - transfer_generation: Option<&'a str>, } } @@ -129,20 +114,6 @@ impl<'a> From<&'a GenerateRequest> for TokenizedGenerateReqInput<'a> { decode_tp_size: req.decode_tp_size, routed_dp_rank: req.routed_dp_rank, disagg_prefill_dp_rank: req.disagg_prefill_dp_rank, - routing_key: (), - require_reasoning: false, - priority: (), - extra_key: (), - no_logs: false, - return_bytes: false, - return_entropy: false, - need_wait_for_mm_inputs: (), - num_items_assigned: (), - encoder_urls: (), - multi_item_delimiter_indices: (), - time_stats: (), - cache_salt: (), - transfer_generation: req.transfer_generation.as_deref(), } } } @@ -209,9 +180,9 @@ mod tests { let bytes = TokenizedGenerateReqInput::from(&req).encode().unwrap(); let val = rmpv::decode::read_value(&mut &bytes[..]).unwrap(); let arr = val.as_array().expect("array"); - // msgspec requires >= 14 (through `stream`); generation is an - // append-only field at 45, after every pre-existing optional field. - assert_eq!(arr.len(), 46, "header ends at transfer_generation"); + // msgspec requires >= 14 (through `stream`); we emit 32 (through + // `disagg_prefill_dp_rank`). Trailing defaulted fields are omitted. + assert_eq!(arr.len(), 32, "header ends at disagg_prefill_dp_rank"); assert_eq!(arr[0].as_str(), Some("TokenizedGenerateReqInput")); assert_eq!(arr[1].as_str(), Some("r1")); assert!(arr[5].is_nil(), "idx 5 must be input_embeds (nil)"); @@ -239,7 +210,6 @@ mod tests { Some(true), "return_hidden_states at idx 16" ); - assert!(arr[45].is_nil(), "generation is nil without a PD room"); } /// The PD block must land on Python's wire indices 25–31, with the filler @@ -257,7 +227,6 @@ mod tests { decode_tp_size: Some(2), routed_dp_rank: Some(3), disagg_prefill_dp_rank: Some(4), - transfer_generation: Some("11111111111111111111111111111111".into()), ..Default::default() }; let bytes = TokenizedGenerateReqInput::from(&req).encode().unwrap(); @@ -276,16 +245,5 @@ mod tests { assert_eq!(arr[29].as_i64(), Some(2), "decode_tp_size at 29"); assert_eq!(arr[30].as_i64(), Some(3), "routed_dp_rank at 30"); assert_eq!(arr[31].as_i64(), Some(4), "disagg_prefill_dp_rank at 31"); - for i in [32, 34, 35, 39, 40, 41, 42, 43, 44] { - assert!(arr[i].is_nil(), "defaulted optional field at idx {i}"); - } - for i in [33, 36, 37, 38] { - assert_eq!(arr[i].as_bool(), Some(false), "false default at idx {i}"); - } - assert_eq!( - arr[45].as_str(), - Some("11111111111111111111111111111111"), - "transfer_generation at append-only idx 45" - ); } } diff --git a/rust/sglang-server/src/message/request.rs b/rust/sglang-server/src/message/request.rs index 1b0d22316d21..10df74b94108 100644 --- a/rust/sglang-server/src/message/request.rs +++ b/rust/sglang-server/src/message/request.rs @@ -7,7 +7,6 @@ use std::sync::LazyLock; use bytes::Bytes; use itertools::izip; use serde::Deserialize; -use uuid::Uuid; use super::io_struct::{ControlRequest, TokenizedGenerateReqInput}; use super::response::ResponseSink; @@ -31,20 +30,6 @@ use crate::utils::{environ::env_i64, error::Error}; static MAX_BATCH_REQS_PER_HTTP_REQ: LazyLock = LazyLock::new(|| env_i64("SGLANG_MAX_BATCH_REQS_PER_HTTP_REQ", 4096)); -/// Mint a full-width, OS-seeded 128-bit lease identity without adding another -/// random-number dependency. UUID v4 supplies 122 random bits and overwrites -/// six RFC version/variant bits; fill those exact positions from independent -/// random bytes of a second UUID before rendering canonical lowercase hex. -fn new_transfer_generation() -> String { - let primary = Uuid::new_v4(); - let refill = Uuid::new_v4(); - let mut bytes = *primary.as_bytes(); - let refill_bytes = refill.as_bytes(); - bytes[6] = (bytes[6] & 0x0f) | (refill_bytes[0] & 0xf0); - bytes[8] = (bytes[8] & 0x3f) | (refill_bytes[1] & 0xc0); - format!("{:032x}", u128::from_be_bytes(bytes)) -} - fn batch_size_exceeds_limit(batch_size: usize, limit: i64) -> bool { limit >= 0 && batch_size as u128 > limit as u128 } @@ -429,9 +414,6 @@ impl GenerateBody { return_text_in_logprobs, bootstrap_host, bootstrap_port, - transfer_generation: bootstrap_room - .as_ref() - .map(|_| new_transfer_generation()), bootstrap_room, bootstrap_pair_key, decode_tp_size, @@ -667,8 +649,6 @@ pub struct GenerateRequest { /// so these are pure passthrough for the scheduler/LB protocol. pub routed_dp_rank: Option, pub disagg_prefill_dp_rank: Option, - /// Internal random room-lease identity, never accepted from the client. - pub transfer_generation: Option, /// Multimodal inputs, carried opaquely. Consumed by the Encoding stage, /// which ships them to the MM worker pool; never read by the tokenizer or /// serialized onto the scheduler header. Boxed so the common text-only @@ -1311,16 +1291,7 @@ mod tests { assert_eq!(p.bootstrap_port, Some(8998)); assert_eq!(p.bootstrap_room, Some(7 + i as i64)); assert_eq!(p.routed_dp_rank, Some(1)); - let generation = p.transfer_generation.as_deref().unwrap(); - assert_eq!(generation.len(), 32); - assert!( - generation - .bytes() - .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)), - "generation must be canonical lowercase 128-bit hex" - ); } - assert_ne!(ps[0].transfer_generation, ps[1].transfer_generation); let (ps, _) = requests( r#"{"text": ["a", "b"], "bootstrap_host": ["h1", "h2"], @@ -1331,7 +1302,6 @@ mod tests { assert_eq!(ps[1].bootstrap_host.as_deref(), Some("h2")); assert_eq!(ps[0].bootstrap_room, Some(10)); assert_eq!(ps[1].bootstrap_room, Some(20)); - assert_ne!(ps[0].transfer_generation, ps[1].transfer_generation); let err = requests(r#"{"text": ["a", "b"], "bootstrap_room": [1, 2, 3]}"#).unwrap_err(); assert!(err.to_string().contains("bootstrap_room"), "{err}"); 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 cbc0fe720090..0a8936f219d6 100644 --- a/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py +++ b/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py @@ -3,24 +3,18 @@ When a decode request is aborted while its prefill->decode KV transfer may still be in flight, the decode side holds its KV pages / req-slot instead of freeing them immediately (which could let the still-in-flight write land on pages already -reused by another request). The pages are released only once every exact prefill -rank for the immutable room generation acks that its transfer drained. A timeout -is diagnostic and never authorizes reuse. +reused by another request). The pages are released once every prefill rank acks +that its transfer drained (CommonKVManager.is_abort_release_safe), or a timeout +fires. See DecodeTransferQueue.resolve_deferred_releases. """ -import os -import threading import unittest from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import patch from sglang.srt.disaggregation import decode as decode_mod -from sglang.srt.disaggregation.common.conn import ( - CommonKVManager, - next_transfer_generation, -) -from sglang.srt.disaggregation.decode import DecodePreallocQueue, DecodeTransferQueue -from sglang.srt.disaggregation.utils import TransferBackend +from sglang.srt.disaggregation.common.conn import CommonKVManager +from sglang.srt.disaggregation.decode import DecodeTransferQueue from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -32,91 +26,91 @@ def _make_manager(): touch (avoids the heavy real __init__).""" mgr = CommonKVManager.__new__(CommonKVManager) mgr._deferred_abort_ack_tracker = {} - mgr._abort_state_lock = threading.Lock() return mgr -GEN_A = "11" * 16 -GEN_B = "22" * 16 - - 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, GEN_A, frozenset({4, 9})) - self.assertFalse(mgr.is_abort_release_safe(room, GEN_A)) + mgr.register_deferred_abort_room(room) + self.assertFalse(mgr.is_abort_release_safe(room, required_acks=2)) - self.assertTrue(mgr.note_abort_ack(room, GEN_A, 4)) - self.assertFalse(mgr.is_abort_release_safe(room, GEN_A)) + mgr.note_abort_ack(room, 0) + self.assertFalse(mgr.is_abort_release_safe(room, required_acks=2)) - self.assertTrue(mgr.note_abort_ack(room, GEN_A, 9)) - self.assertTrue(mgr.is_abort_release_safe(room, GEN_A)) + mgr.note_abort_ack(room, 1) + 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, GEN_A, frozenset({0, 1})) - mgr.note_abort_ack(room, GEN_A, 0) - mgr.note_abort_ack(room, GEN_A, 0) - self.assertFalse(mgr.is_abort_release_safe(room, GEN_A)) + mgr.register_deferred_abort_room(room) + mgr.note_abort_ack(room, 0) + mgr.note_abort_ack(room, 0) # 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_invalid_and_unexpected_rank_identities_are_rejected(self): + def test_single_rank_fast_path(self): mgr = _make_manager() room = 102 - mgr.register_deferred_abort_room(room, GEN_A, frozenset({7})) - for bad_rank in (True, -1, 8, 2**63): - self.assertFalse(mgr.note_abort_ack(room, GEN_A, bad_rank)) - self.assertFalse(mgr.is_abort_release_safe(room, GEN_A)) - self.assertTrue(mgr.note_abort_ack(room, GEN_A, 7)) - self.assertTrue(mgr.is_abort_release_safe(room, GEN_A)) + 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, GEN_A, frozenset({0})) - mgr.note_abort_ack(room, GEN_A, 0) - mgr.clear_deferred_abort_state(room, GEN_A) - self.assertNotIn((room, GEN_A), mgr._deferred_abort_ack_tracker) - self.assertFalse(mgr.is_abort_release_safe(room, GEN_A)) + mgr.register_deferred_abort_room(room) + mgr.note_abort_ack(room, 0) + 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)) def test_ack_before_register_is_dropped(self): + # An ack for a room that isn't actively held must not be recorded (it + # would otherwise pollute a later request reusing the same room). mgr = _make_manager() room = 104 - self.assertFalse(mgr.note_abort_ack(room, GEN_A, 0)) - self.assertNotIn((room, GEN_A), mgr._deferred_abort_ack_tracker) - - def test_room_reuse_is_generation_isolated(self): + mgr.note_abort_ack(room, 0) # 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 - expected = frozenset({0, 1}) - mgr.register_deferred_abort_room(room, GEN_A, expected) - mgr.note_abort_ack(room, GEN_A, 0) - mgr.clear_deferred_abort_state(room, GEN_A) - mgr.register_deferred_abort_room(room, GEN_B, expected) - - self.assertFalse(mgr.note_abort_ack(room, GEN_A, 1)) - mgr.note_abort_ack(room, GEN_B, 0) - self.assertFalse(mgr.is_abort_release_safe(room, GEN_B)) - mgr.note_abort_ack(room, GEN_B, 1) - self.assertTrue(mgr.is_abort_release_safe(room, GEN_B)) - - def test_rearming_is_idempotent_but_expected_set_is_immutable(self): + + # 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 - expected = frozenset({0, 1}) - mgr.register_deferred_abort_room(room, GEN_A, expected) - mgr.note_abort_ack(room, GEN_A, 0) - mgr.register_deferred_abort_room(room, GEN_A, expected) - self.assertFalse(mgr.is_abort_release_safe(room, GEN_A)) - with self.assertRaises(RuntimeError): - mgr.register_deferred_abort_room(room, GEN_A, frozenset({0})) - - def test_expected_rank_set_cannot_be_empty_or_malformed(self): - mgr = _make_manager() - for bad in (frozenset(), frozenset({-1}), frozenset({True})): - with self.assertRaises(ValueError): - mgr.register_deferred_abort_room(107, GEN_A, bad) + mgr.register_deferred_abort_room(room) + mgr.note_abort_ack(room, 0) + mgr.note_abort_ack(room, 1) + 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) + self.assertFalse(mgr.is_abort_release_safe(room, required_acks=2)) class _FakeIdxAllocator: @@ -139,30 +133,16 @@ def _make_queue(timeout=30.0): return q -def _make_decode_req(room, idx, mgr, ranks=(0,), generation=GEN_A): +def _make_decode_req(room, idx, mgr, n_prefill_ranks=1): receiver = SimpleNamespace( kv_mgr=mgr, - bootstrap_infos=[{"prefill_unique_rank": rank} for rank in ranks], - transfer_generation=generation, - abort_notified=False, - clear=MagicMock(), - retry_abort_notification=MagicMock(), - failure_exception=MagicMock(side_effect=RuntimeError("transfer failed")), + # One entry per prefill rank the decode notified of the abort; its length + # is the required drain-ack count (see DecodeTransferQueue._defer_release). + bootstrap_infos=[{"rank": r} for r in range(n_prefill_ranks)], + clear=lambda: None, ) - - def arm_abort_intent(): - if receiver.abort_notified: - return - mgr.register_deferred_abort_room(room, generation, frozenset(ranks)) - receiver.abort_notified = True - - receiver.arm_abort_intent = arm_abort_intent return SimpleNamespace( - req=SimpleNamespace( - bootstrap_room=room, - rid=f"request-{room}", - return_logprob=False, - ), + req=SimpleNamespace(bootstrap_room=room), kv_receiver=receiver, metadata_buffer_index=idx, ) @@ -179,7 +159,10 @@ def test_holds_until_drained_then_releases(self): mgr = _make_manager() room, idx = 200, 7 q = _make_queue() - dreq = _make_decode_req(room, idx, mgr, ranks=(4, 9)) + dreq = _make_decode_req(room, idx, mgr, n_prefill_ranks=2) + # In production the room is armed in abort_request when the ABORT is + # sent, before the scheduler defers here. + mgr.register_deferred_abort_room(room) q._defer_release(dreq) with patch.object(decode_mod, "release_kv_cache") as rel: @@ -189,13 +172,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, GEN_A, 4) + mgr.note_abort_ack(room, 0) 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, GEN_A, 9) + mgr.note_abort_ack(room, 1) q.resolve_deferred_releases() rel.assert_called_once_with(dreq.req, q.tree_cache, is_insert=False) @@ -203,36 +186,35 @@ def test_holds_until_drained_then_releases(self): self.assertEqual(q._deferred_releases, []) self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [idx]) self.assertEqual(q.metadata_buffers.bootstrap_room[idx], 0) - self.assertNotIn((room, GEN_A), mgr._deferred_abort_ack_tracker) + self.assertNotIn(room, mgr._deferred_abort_ack_tracker) self.assertIsNone(dreq.kv_receiver) - def test_timeout_only_diagnoses_retries_and_retains_quarantine(self): + def test_releases_on_timeout_without_ack(self): mgr = _make_manager() room, idx = 300, 3 q = _make_queue(timeout=30.0) - dreq = _make_decode_req(room, idx, mgr) - q._defer_release(dreq) - q._deferred_releases[0].next_diagnostic_at = float("-inf") + dreq = _make_decode_req(room, idx, mgr, n_prefill_ranks=1) + # Force an already-expired deadline (no ack will ever arrive). + q._deferred_releases.append((dreq, float("-inf"), idx, 1)) with patch.object(decode_mod, "release_kv_cache") as rel: q.resolve_deferred_releases() - rel.assert_not_called() + rel.assert_called_once_with(dreq.req, q.tree_cache, is_insert=False) - self.assertEqual(len(q._deferred_releases), 1) - self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, []) - self.assertIs(dreq.kv_receiver, q._deferred_releases[0].kv_receiver) - dreq.kv_receiver.retry_abort_notification.assert_called_once_with() - self.assertEqual(q._deferred_releases[0].diagnostic_count, 1) + self.assertEqual(q._deferred_releases, []) + self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [idx]) + self.assertIsNone(dreq.kv_receiver) - def test_failed_release_is_fail_stop_while_other_owner_completes(self): + def test_failed_release_is_isolated_and_not_retried(self): + # A raising _do_release must drop the entry (no double-free on retry) and + # not brick resolve for the remaining entries or subsequent calls. mgr = _make_manager() q = _make_queue() - bad = _make_decode_req(701, 2, mgr, generation=GEN_A) - good = _make_decode_req(700, 1, mgr, generation=GEN_B) - q._defer_release(bad) - q._defer_release(good) - mgr.note_abort_ack(701, GEN_A, 0) - mgr.note_abort_ack(700, GEN_B, 0) + good = _make_decode_req(700, 1, mgr) + bad = _make_decode_req(701, 2, mgr) + # Both already past deadline -> both selected for release. + q._deferred_releases.append((bad, float("-inf"), 2, 1)) + q._deferred_releases.append((good, float("-inf"), 1, 1)) calls = [] @@ -242,97 +224,12 @@ def fake_release(req, tree_cache, is_insert): raise RuntimeError("boom") with patch.object(decode_mod, "release_kv_cache", side_effect=fake_release): - q.resolve_deferred_releases() + q.resolve_deferred_releases() # must not raise + # The good one still released despite the bad one throwing. self.assertIn(good.req, calls) - self.assertEqual(len(q._deferred_releases), 1) - failed = q._deferred_releases[0] - self.assertIs(failed.decode_req, bad) - self.assertEqual(failed.release_phase, "KV_RELEASE_ENTERED") - self.assertIsNotNone(failed.release_error) - prior_calls = list(calls) + # Nothing left held, and a second call is a clean no-op (no retry). + self.assertEqual(q._deferred_releases, []) q.resolve_deferred_releases() - self.assertEqual(calls, prior_calls) - - def test_lost_return_from_non_idempotent_free_is_never_retried(self): - mgr = _make_manager() - q = _make_queue() - dreq = _make_decode_req(702, 5, mgr) - q._defer_release(dreq) - mgr.note_abort_ack(702, GEN_A, 0) - - real_free = q.req_to_metadata_buffer_idx_allocator.free - - def committed_then_interrupted(idx): - real_free(idx) - raise KeyboardInterrupt("lost return") - - q.req_to_metadata_buffer_idx_allocator.free = MagicMock( - side_effect=committed_then_interrupted - ) - with self.assertRaises(KeyboardInterrupt): - q.resolve_deferred_releases() - self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [5]) - self.assertEqual( - q._deferred_releases[0].release_phase, - "METADATA_INDEX_FREE_ENTERED", - ) - q.resolve_deferred_releases() - q.req_to_metadata_buffer_idx_allocator.free.assert_called_once_with(5) - - def test_explicit_false_environment_cannot_authorize_timeout_release(self): - mgr = _make_manager() - with patch.dict( - os.environ, - {"SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE": "false"}, - ): - q = DecodeTransferQueue( - gloo_group=None, - req_to_metadata_buffer_idx_allocator=_FakeIdxAllocator(), - tp_rank=0, - metadata_buffers=SimpleNamespace(bootstrap_room={}), - scheduler=SimpleNamespace(spec_algorithm=None), - tree_cache=object(), - ) - self.assertTrue(q.enable_deferred_kv_release) - dreq = _make_decode_req(703, 6, mgr) - q._defer_release(dreq) - q._deferred_releases[0].next_diagnostic_at = float("-inf") - with patch.object(decode_mod, "release_kv_cache") as rel: - q.resolve_deferred_releases() - rel.assert_not_called() - self.assertEqual(len(q._deferred_releases), 1) - - def test_false_compatibility_fields_cannot_bypass_failure_quarantine(self): - mgr = _make_manager() - mgr.enable_deferred_decode_kv_release = False - q = _make_queue() - q.enable_deferred_kv_release = False - q.enable_staging = False - q.gloo_group = None - q.tp_rank = 0 - q.queue = [_make_decode_req(704, 8, mgr)] - q.queue[0].hicache_restore_status = None - # This is specifically a decode-initiated failure with a possible - # remote writer. A propagated prefill failure has already stopped its - # writer and intentionally uses the immediate-safe release path. - q.queue[0].kv_receiver.arm_abort_intent() - q._poll_with_metadata_gate = MagicMock(return_value=[decode_mod.KVPoll.Failed]) - q._clean_hicache_prefetch_resources = MagicMock() - q.scheduler = SimpleNamespace( - enable_decode_hicache=False, - enable_hisparse=False, - output_streamer=SimpleNamespace(stream_output=MagicMock()), - metrics_reporter=SimpleNamespace(enable_metrics=False), - ) - - with patch.object(decode_mod, "release_kv_cache") as release: - q.pop_transferred() - - release.assert_not_called() - self.assertEqual(q.queue, []) - self.assertEqual(len(q._deferred_releases), 1) - self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, []) - self.assertIsNotNone(q._deferred_releases[0].decode_req.kv_receiver) def test_defer_release_records_deadline_and_idx(self): mgr = _make_manager() @@ -340,73 +237,10 @@ def test_defer_release_records_deadline_and_idx(self): dreq = _make_decode_req(room=400, idx=9, mgr=mgr) q._defer_release(dreq) self.assertEqual(len(q._deferred_releases), 1) - held = q._deferred_releases[0] - self.assertIs(held.decode_req, dreq) - self.assertEqual(held.metadata_buffer_index, 9) - self.assertEqual(held.transfer_generation, GEN_A) - self.assertIsInstance(held.next_diagnostic_at, float) - - -class TestTransferGenerationPlumbing(CustomTestCase): - @staticmethod - def _make_prealloc_queue(): - queue = DecodePreallocQueue.__new__(DecodePreallocQueue) - queue.transfer_backend = TransferBackend.NIXL - queue.scheduler = SimpleNamespace( - output_streamer=SimpleNamespace(stream_output=MagicMock()) - ) - queue.pending_reqs = [] - queue.retracted_queue = [] - queue._check_if_req_exceed_kv_capacity = MagicMock(return_value=False) - queue._resolve_prefill_dp_rank = MagicMock(return_value=0) - receiver = SimpleNamespace(init=MagicMock()) - queue._create_receiver_and_enqueue = MagicMock( - return_value=SimpleNamespace(kv_receiver=receiver) - ) - return queue, receiver - - def test_rebootstrap_advances_identically_on_each_decode_rank(self): - advanced = [] - for _ in range(2): - queue, receiver = self._make_prealloc_queue() - req = SimpleNamespace( - bootstrap_host="prefill.internal", - transfer_generation=GEN_A, - return_logprob=False, - ) - - queue.add(req, is_rebootstrap=True) - - advanced.append(req.transfer_generation) - queue._create_receiver_and_enqueue.assert_called_once_with( - req, is_rebootstrap=True - ) - receiver.init.assert_called_once_with(0) - - self.assertEqual(advanced, [next_transfer_generation(GEN_A)] * 2) - self.assertNotIn(advanced[0], (GEN_A, GEN_B)) - self.assertEqual(len(advanced[0]), 32) - - def test_missing_generation_fails_request_before_receiver_construction(self): - queue, _ = self._make_prealloc_queue() - req = SimpleNamespace( - bootstrap_host="prefill.internal", - transfer_generation=None, - return_logprob=False, - ) - - with patch.object(decode_mod, "prepare_abort") as prepare_abort: - queue.add(req) - - queue._create_receiver_and_enqueue.assert_not_called() - prepare_abort.assert_called_once() - self.assertEqual( - prepare_abort.call_args.kwargs["status_code"], - decode_mod.HTTPStatus.INTERNAL_SERVER_ERROR, - ) - queue.scheduler.output_streamer.stream_output.assert_called_once_with( - [req], False - ) + held_req, deadline, held_idx, required = q._deferred_releases[0] + self.assertIs(held_req, dreq) + self.assertEqual(held_idx, 9) + self.assertIsInstance(deadline, float) if __name__ == "__main__": diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 39d5d17373d9..a2bdf5aabdb8 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -14,8 +14,6 @@ ) from sglang.srt.disaggregation.common.staging_handler import ( DecodeStagingHandler, - STAGING_REQ_WIRE_TAG, - STAGING_RSP_WIRE_TAG, handle_staging_req, ) from sglang.srt.disaggregation.common.utils import ( @@ -260,11 +258,9 @@ def test_new_watermark_subscriber_receives_current_allocator_state(self): class TestMooncakePPStaging(unittest.TestCase): def test_staging_response_targets_requesting_pp_rank(self): - generation = "0123456789abcdef0123456789abcdef" sock = Mock() receiver = SimpleNamespace( chunk_staging_infos=[], - transfer_generation=generation, _connect_to_bootstrap_server=Mock(return_value=(sock, threading.Lock())), ) allocator = SimpleNamespace( @@ -279,15 +275,7 @@ def test_staging_response_targets_requesting_pp_rank(self): target = {"pp_rank": 3} handle_staging_req( - [ - STAGING_REQ_WIRE_TAG, - generation.encode("ascii"), - b"7", - b"0", - b"1", - b"peer", - b"3", - ], + [b"STAGING_REQ", b"7", b"0", b"1", b"peer", b"3"], allocator, kv_args, attn_tp_size=16, @@ -299,9 +287,6 @@ def test_staging_response_targets_requesting_pp_rank(self): receiver._connect_to_bootstrap_server.assert_called_once_with(target) sock.send_multipart.assert_called_once() - response = sock.send_multipart.call_args.args[0] - self.assertEqual(response[0], STAGING_RSP_WIRE_TAG) - self.assertEqual(response[1], generation.encode("ascii")) @patch( "sglang.srt.disaggregation.common.staging_buffer.gather_all_layers_to_staging" diff --git a/test/registered/unit/disaggregation/test_mooncake_generation_control.py b/test/registered/unit/disaggregation/test_mooncake_generation_control.py deleted file mode 100644 index 4180863552e2..000000000000 --- a/test/registered/unit/disaggregation/test_mooncake_generation_control.py +++ /dev/null @@ -1,192 +0,0 @@ -"""Adversarial Mooncake control-plane lease identity tests.""" - -import struct -import threading -from collections import defaultdict -from types import SimpleNamespace -from unittest.mock import MagicMock, patch - -from sglang.srt.disaggregation.base.conn import KVPoll -from sglang.srt.disaggregation.common.conn import CommonKVReceiver -from sglang.srt.disaggregation.common.staging_handler import DecodeStagingHandler -from sglang.srt.disaggregation.mooncake.conn import ( - MOONCAKE_AUX_DATA_WIRE_TAG, - MOONCAKE_CHUNK_READY_WIRE_TAG, - MOONCAKE_TRANSFER_STATUS_WIRE_TAG, - MooncakeKVManager, - MooncakeKVReceiver, -) -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=1, suite="base-a-test-cpu") - -GEN_A = "a1" * 16 -GEN_B = "b2" * 16 -ROOM = 7 -SESSION = "decode-session" - - -def _manager(generation=GEN_B, expected=frozenset({3, 5})): - manager = object.__new__(MooncakeKVManager) - manager._abort_state_lock = threading.Lock() - manager._receive_transfer_leases = {} - manager.prefill_response_tracker = defaultdict(set) - manager.request_status = {ROOM: KVPoll.WaitingForInput} - manager.required_prefill_response_num_table = {ROOM: len(expected)} - manager.failure_records = {} - manager.failure_lock = threading.Lock() - manager.enable_staging = True - manager._staging_handler = MagicMock() - manager.kv_args = SimpleNamespace() - manager._register_receive_transfer_lease(ROOM, generation, SESSION, expected) - return manager - - -def _chunk(generation, rank=3, session=SESSION): - return [ - MOONCAKE_CHUNK_READY_WIRE_TAG, - generation.encode("ascii"), - str(ROOM).encode("ascii"), - b"0", - b"4", - b"2", - session.encode("ascii"), - str(rank).encode("ascii"), - ] - - -def _aux(generation, rank=3, payload=b"data"): - return [ - MOONCAKE_AUX_DATA_WIRE_TAG, - generation.encode("ascii"), - str(ROOM).encode("ascii"), - str(rank).encode("ascii"), - b"0", - b"1", - struct.pack(">I", len(payload)), - payload, - ] - - -def _status(generation, rank, status=KVPoll.Success): - return [ - MOONCAKE_TRANSFER_STATUS_WIRE_TAG, - generation.encode("ascii"), - str(ROOM).encode("ascii"), - str(status).encode("ascii"), - str(rank).encode("ascii"), - ] - - -class TestMooncakeGenerationControl(CustomTestCase): - def test_receiver_init_binds_immutable_exact_rank_set(self): - manager = _manager() - manager._clear_receive_transfer_lease(ROOM, GEN_B) - manager.addr_to_rooms_tracker = defaultdict(set) - manager.get_session_id = MagicMock(return_value=SESSION) - - receiver = MooncakeKVReceiver( - manager, - "prefill:8998", - ROOM, - transfer_generation=GEN_B, - ) - receiver.bootstrap_infos = [ - {"prefill_unique_rank": 3, "is_dummy": False}, - {"prefill_unique_rank": 4, "is_dummy": True}, - {"prefill_unique_rank": 5, "is_dummy": False}, - ] - with patch.object(CommonKVReceiver, "init", return_value=None): - receiver.init(prefill_dp_rank=0) - - lease = manager._receive_transfer_leases[ROOM] - self.assertEqual(lease.generation, GEN_B) - self.assertEqual(lease.session_id, SESSION) - self.assertEqual(lease.expected_prefill_ranks, frozenset({3, 5})) - with self.assertRaisesRegex(RuntimeError, "immutable"): - manager._register_receive_transfer_lease( - ROOM, GEN_B, SESSION, frozenset({3, 6}) - ) - - def test_old_chunk_cannot_scatter_reused_room(self): - manager = _manager() - - self.assertFalse(manager._handle_chunk_ready(_chunk(GEN_A))) - self.assertFalse(manager._handle_chunk_ready(_chunk(GEN_B, rank=99))) - self.assertFalse(manager._handle_chunk_ready(_chunk(GEN_B, session="old"))) - self.assertFalse( - manager._handle_chunk_ready( - [b"CHUNK_READY", b"7", b"0", b"4", b"2", SESSION.encode()] - ) - ) - manager._staging_handler.handle_chunk_arrived.assert_not_called() - - self.assertTrue(manager._handle_chunk_ready(_chunk(GEN_B, rank=3))) - manager._staging_handler.handle_chunk_arrived.assert_called_once_with( - ROOM, 0, 4, 2, 3 - ) - - def test_old_aux_cannot_write_reused_room(self): - manager = _manager() - target = "sglang.srt.disaggregation.mooncake.conn.AuxDataCodec.deserialize_data_to_buffer" - with patch(target) as deserialize: - self.assertFalse(manager._handle_aux_data(_aux(GEN_A))) - self.assertFalse(manager._handle_aux_data(_aux(GEN_B, rank=99))) - self.assertFalse( - manager._handle_aux_data( - [b"AUX_DATA", b"7", b"0", b"1", struct.pack(">I", 4), b"data"] - ) - ) - deserialize.assert_not_called() - - self.assertTrue(manager._handle_aux_data(_aux(GEN_B, rank=3))) - deserialize.assert_called_once_with(manager.kv_args, 0, 1, b"data") - - def test_status_requires_exact_generation_and_all_distinct_expected_ranks(self): - manager = _manager() - - self.assertFalse(manager._handle_transfer_status(_status(GEN_A, 3))) - self.assertFalse(manager._handle_transfer_status(_status(GEN_B, 99))) - self.assertFalse(manager._handle_transfer_status([b"7", b"4", b"3"])) - self.assertEqual(manager.request_status[ROOM], KVPoll.WaitingForInput) - - self.assertTrue(manager._handle_transfer_status(_status(GEN_B, 3))) - self.assertTrue(manager._handle_transfer_status(_status(GEN_B, 3))) - self.assertEqual(manager.request_status[ROOM], KVPoll.WaitingForInput) - self.assertEqual(manager.prefill_response_tracker[ROOM], {3}) - - self.assertTrue(manager._handle_transfer_status(_status(GEN_B, 5))) - self.assertEqual(manager.prefill_response_tracker[ROOM], {3, 5}) - self.assertEqual(manager.request_status[ROOM], KVPoll.Success) - - def test_delayed_old_receiver_clear_cannot_erase_new_lease(self): - manager = _manager(generation=GEN_A) - self.assertTrue(manager._clear_receive_transfer_lease(ROOM, GEN_A)) - manager.request_status[ROOM] = KVPoll.WaitingForInput - manager.required_prefill_response_num_table[ROOM] = 2 - # The old receiver has already cleared its identity. A duplicate clear - # must not treat that absence as authority over state B is publishing. - self.assertFalse(manager._clear_receive_transfer_lease(ROOM, GEN_A)) - self.assertEqual(manager.request_status[ROOM], KVPoll.WaitingForInput) - manager._register_receive_transfer_lease( - ROOM, GEN_B, SESSION, frozenset({3, 5}) - ) - - self.assertFalse(manager._clear_receive_transfer_lease(ROOM, GEN_A)) - self.assertEqual(manager._receive_transfer_leases[ROOM].generation, GEN_B) - self.assertEqual(manager.request_status[ROOM], KVPoll.WaitingForInput) - - def test_staging_fan_in_deduplicates_exact_writer_rank(self): - handler = object.__new__(DecodeStagingHandler) - handler._room_to_receiver = {ROOM: object()} - handler._writer_counts = {} - handler.num_writers_for = MagicMock(return_value=2) - handler.submit_chunk_scatter = MagicMock(return_value=True) - - self.assertFalse(handler.handle_chunk_arrived(ROOM, 0, 4, 2, 3)) - self.assertFalse(handler.handle_chunk_arrived(ROOM, 0, 4, 2, 3)) - handler.submit_chunk_scatter.assert_not_called() - - self.assertTrue(handler.handle_chunk_arrived(ROOM, 0, 4, 2, 5)) - handler.submit_chunk_scatter.assert_called_once_with(ROOM, 0, 4, 2) diff --git a/test/registered/unit/disaggregation/test_nixl_backend_basic.py b/test/registered/unit/disaggregation/test_nixl_backend_basic.py index e6ab6c0c3c11..2bc8a210738a 100644 --- a/test/registered/unit/disaggregation/test_nixl_backend_basic.py +++ b/test/registered/unit/disaggregation/test_nixl_backend_basic.py @@ -6,7 +6,6 @@ import types import unittest from collections import defaultdict -from dataclasses import replace from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -14,13 +13,7 @@ from sglang.srt.disaggregation.base.conn import KVPoll from sglang.srt.disaggregation.common.conn import CommonKVManager -from sglang.srt.disaggregation.common.staging_handler import ( - STAGING_REQ_WIRE_TAG, - STAGING_RSP_WIRE_TAG, - PrefillStagingContext, - handle_staging_req, - handle_staging_rsp, -) +from sglang.srt.disaggregation.common.staging_handler import PrefillStagingContext from sglang.srt.disaggregation.common.utils import pack_int_lists from sglang.srt.disaggregation.nixl.conn import ( KVArgsRegisterInfo, @@ -30,21 +23,12 @@ TransferInfo, TransferKVChunk, TransferStatus, - repeat_indices_over_layers, ) -from sglang.srt.disaggregation.utils import DisaggregationMode 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") -TRANSFER_GENERATION = "11" * 16 -OLD_TRANSFER_GENERATION = "22" * 16 - - -def completion_notification(room, suffix, generation=TRANSFER_GENERATION): - return f"nixlv2_{generation}_{room}_{suffix}" - class NotificationFakeAgent: def __init__(self, messages): @@ -136,8 +120,6 @@ def test_from_zmq_parses_required_fields(self): b"2", pack_int_lists(state_indices, "i"), b"11", - b"0", - TRANSFER_GENERATION.encode("ascii"), ] info = TransferInfo.from_zmq(msg) @@ -151,7 +133,6 @@ def test_from_zmq_parses_required_fields(self): self.assertEqual(info.required_dst_info_num, 2) self.assertEqual(info.dst_state_indices, state_indices) self.assertEqual(info.decode_prefix_len, 11) - self.assertEqual(info.transfer_generation, TRANSFER_GENERATION) def test_from_zmq_defaults_optional_fields(self): info = TransferInfo.from_zmq( @@ -163,10 +144,6 @@ def test_from_zmq_defaults_optional_fields(self): np.array([1], dtype=np.int32).tobytes(), b"0", b"1", - b"", - b"", - b"", - TRANSFER_GENERATION.encode("ascii"), ] ) @@ -185,8 +162,6 @@ def test_decode_radix_full_hit_is_not_dummy(self): b"1", b"", b"128", - b"", - TRANSFER_GENERATION.encode("ascii"), ] ) @@ -204,8 +179,6 @@ def test_empty_indices_without_decode_prefix_is_dummy(self): b"1", b"", b"0", - b"", - TRANSFER_GENERATION.encode("ascii"), ] ) @@ -303,20 +276,8 @@ def test_from_zmq_allows_missing_state_and_staging_fields(self): class TestNixlTransferStatus(CustomTestCase): - def test_completion_source_rank_set_is_exact_and_immutable(self): - status = TransferStatus(transfer_generation=TRANSFER_GENERATION) - status.bind_expected_source_ranks(frozenset({4, 9})) - - self.assertTrue(status.accepts_source_rank(4)) - for unexpected in (True, -1, 0, 10): - self.assertFalse(status.accepts_source_rank(unexpected)) - self.assertEqual(status.num_pp_ranks_expected, 2) - status.bind_expected_source_ranks(frozenset({4, 9})) - with self.assertRaises(RuntimeError): - status.bind_expected_source_ranks(frozenset({4})) - def test_not_done_until_aux_and_expected_count_arrive(self): - status = TransferStatus(transfer_generation=TRANSFER_GENERATION) + status = TransferStatus() self.assertFalse(status.is_done()) @@ -333,7 +294,7 @@ def test_not_done_until_aux_and_expected_count_arrive(self): self.assertTrue(status.is_done()) def test_zero_kv_aux_only_completion(self): - status = TransferStatus(transfer_generation=TRANSFER_GENERATION) + status = TransferStatus() status.received_aux = True status.num_pp_ranks_expected = 1 status.expected_kvs_per_pp[0] = 0 @@ -341,7 +302,7 @@ def test_zero_kv_aux_only_completion(self): self.assertTrue(status.is_done()) def test_multi_pp_requires_each_rank_expected_chunks(self): - status = TransferStatus(transfer_generation=TRANSFER_GENERATION) + status = TransferStatus() status.received_aux = True status.num_pp_ranks_expected = 2 status.expected_kvs_per_pp[0] = 1 @@ -354,7 +315,7 @@ def test_multi_pp_requires_each_rank_expected_chunks(self): self.assertTrue(status.is_done()) def test_state_required_completion_waits_for_all_pp_ranks(self): - status = TransferStatus(transfer_generation=TRANSFER_GENERATION) + status = TransferStatus() status.received_aux = True status.num_pp_ranks_expected = 2 status.expected_kvs_per_pp[0] = 0 @@ -378,148 +339,6 @@ def test_last_zero_page_chunk_is_sent_for_aux_only_completion(self): self.assertFalse(sender.should_send_kv_chunk(0, last_chunk=False)) self.assertTrue(sender.should_send_kv_chunk(3, last_chunk=False)) - def test_both_data_planes_support_overlap_early_send(self): - sender = object.__new__(NixlKVSender) - - for use_torch_transfer in (False, True): - sender.kv_mgr = SimpleNamespace(_use_torch_transfer=use_torch_transfer) - self.assertTrue(sender.supports_overlap_early_send) - - def test_send_consumes_early_send_event_into_transfer_chunk(self): - sender = object.__new__(NixlKVSender) - wait_event = object() - sender._early_send_wait_event = wait_event - sender._send_failed = False - sender._transfer_start_time = 1.0 - sender.bootstrap_room = 7 - sender.chunk_id = 2 - sender.aux_index = 3 - sender.has_sent = False - sender.kv_mgr = SimpleNamespace(add_transfer_request=MagicMock()) - prepared_indices = np.array([5], dtype=np.int32) - sender._prepare_send_indices = MagicMock( - return_value=( - prepared_indices, - slice(0, 1), - False, - False, - ) - ) - sender._record_transfer_indices = MagicMock() - - sender.send(np.array([5], dtype=np.int32), num_kv_tokens=4) - - self.assertIsNone(sender._early_send_wait_event) - sender.kv_mgr.add_transfer_request.assert_called_once_with( - 7, - prepared_indices, - slice(0, 1), - False, - 2, - 3, - None, - 4, - wait_event=wait_event, - ) - - -class TestNixlBootstrapErrors(CustomTestCase): - def _make_manager(self): - mgr = object.__new__(NixlKVManager) - mgr.enable_staging = False - mgr._bootstrap_errors = {} - mgr.exceptions = {} - mgr.transfer_infos = {} - mgr.req_to_decode_prefix_len = {} - mgr.prep_handles = {} - mgr._handle_abort_notification = MagicMock(return_value=False) - mgr.record_failure = MagicMock() - mgr.update_status = MagicMock() - return mgr - - def test_registration_error_is_propagated_to_later_room(self): - mgr = self._make_manager() - unsupported = NotImplementedError("unsupported PyTorch transfer topology") - mgr._add_remote_peer = MagicMock(side_effect=unsupported) - registration_message = [ - b"NixlMsgGuard", - b"None", - b"endpoint", - b"1234", - b"decode-agent", - ] - room_message = [ - b"NixlMsgGuard", - b"41", - b"endpoint", - b"1234", - b"decode-agent", - b"", - b"0", - b"1", - b"", - b"", - b"0", - TRANSFER_GENERATION.encode("ascii"), - ] - - with patch.object(KVArgsRegisterInfo, "from_zmq", return_value=object()): - with self.assertRaises(NotImplementedError): - mgr._handle_bootstrap_message(registration_message) - self.assertIs(mgr._bootstrap_errors["decode-agent"], unsupported) - - with self.assertRaises(NotImplementedError): - mgr._handle_bootstrap_message(room_message) - - self.assertIs(mgr.exceptions[41], unsupported) - mgr.record_failure.assert_called_once_with(41, str(unsupported)) - mgr.update_status.assert_called_once_with(41, KVPoll.Failed) - - def test_listener_continues_after_message_error(self): - mgr = self._make_manager() - mgr.server_socket = SimpleNamespace( - recv_multipart=MagicMock(side_effect=[[b"bad"], [b"good"], SystemExit()]) - ) - mgr._handle_bootstrap_message = MagicMock( - side_effect=[ValueError("bad bootstrap"), None] - ) - - with patch( - "sglang.srt.disaggregation.nixl.conn.threading.Thread" - ) as thread_cls: - mgr._start_bootstrap_thread() - target = thread_cls.call_args.kwargs["target"] - - with self.assertRaises(SystemExit): - target() - - self.assertEqual(mgr._handle_bootstrap_message.call_count, 2) - self.assertTrue(thread_cls.call_args.kwargs["daemon"]) - - def test_peer_import_failure_does_not_poison_registration_table(self): - mgr = self._make_manager() - mgr._use_torch_transfer = True - mgr.attn_tp_size = 1 - mgr.disaggregation_mode = object() - mgr.decode_kv_args_table = {} - mgr.requires_dcp_relayout = MagicMock(return_value=False) - mgr.agent = SimpleNamespace( - add_remote_agent=MagicMock(side_effect=RuntimeError("metadata rejected")) - ) - peer = SimpleNamespace( - agent_name="decode-agent", - agent_metadata=b"metadata", - decode_tp_size=1, - dst_dcp_size=1, - dst_dcp_rank=0, - dst_kv_mem_kinds=["VRAM"], - ) - - with self.assertRaisesRegex(RuntimeError, "metadata rejected"): - mgr._add_remote_peer(peer) - - self.assertNotIn("decode-agent", mgr.decode_kv_args_table) - class TestNixlAbortHandling(CustomTestCase): def _make_manager(self, request_status=None): @@ -528,15 +347,9 @@ def _make_manager(self, request_status=None): mgr._connect = MagicMock() mgr.failure_lock = threading.Lock() mgr.failure_records = {} - mgr._abort_state_lock = threading.Lock() - mgr._deferred_ack_targets = {} - mgr._active_transfer_generations = {} - mgr._abort_ack_receipts = set() - mgr._staging_outstanding = {room: 1 for room in mgr.request_status} - mgr._send_abort_ack = MagicMock(return_value=True) + # These cases cover the legacy no-ack behavior; the deferred-release ack + # path is exercised in test_nixl_deferred_kv_release.py. mgr.enable_deferred_decode_kv_release = False - for room in mgr.request_status: - mgr.register_transfer_generation(room, TRANSFER_GENERATION) return mgr def test_given_known_incomplete_room_when_abort_arrives_then_room_fails_without_ack( @@ -545,13 +358,7 @@ def test_given_known_incomplete_room_when_abort_arrives_then_room_fails_without_ mgr = self._make_manager({11: KVPoll.WaitingForInput}) handled = mgr._handle_abort_notification( - [ - b"ABORT", - b"11", - b"127.0.0.1", - b"5555", - TRANSFER_GENERATION.encode("ascii"), - ] + [b"ABORT", b"11", b"127.0.0.1", b"5555"] ) self.assertTrue(handled) @@ -566,13 +373,7 @@ def test_given_successful_room_when_abort_arrives_then_status_is_preserved(self) mgr = self._make_manager({12: KVPoll.Success}) handled = mgr._handle_abort_notification( - [ - b"ABORT", - b"12", - b"127.0.0.1", - b"5556", - TRANSFER_GENERATION.encode("ascii"), - ] + [b"ABORT", b"12", b"127.0.0.1", b"5556"] ) self.assertTrue(handled) @@ -584,13 +385,7 @@ def test_given_unknown_room_when_abort_arrives_then_status_remains_absent(self): mgr = self._make_manager() handled = mgr._handle_abort_notification( - [ - b"ABORT", - b"14", - b"127.0.0.1", - b"5557", - TRANSFER_GENERATION.encode("ascii"), - ] + [b"ABORT", b"14", b"127.0.0.1", b"5557"] ) self.assertTrue(handled) @@ -602,13 +397,7 @@ def test_given_malformed_abort_when_handled_then_no_exception_or_ack(self): mgr = self._make_manager({13: KVPoll.WaitingForInput}) handled = mgr._handle_abort_notification( - [ - b"ABORT", - b"invalid-room", - b"127.0.0.1", - b"5558", - TRANSFER_GENERATION.encode("ascii"), - ] + [b"ABORT", b"invalid-room", b"127.0.0.1", b"5558"] ) self.assertTrue(handled) @@ -657,7 +446,6 @@ def _make_manager(self, room): dst_aux_index=0, required_dst_info_num=1, dst_state_indices=[], - transfer_generation=TRANSFER_GENERATION, ) } } @@ -684,16 +472,11 @@ def _make_manager(self, room): mgr.enable_deferred_decode_kv_release = False mgr._staging_ctx = None mgr._staging_outstanding = defaultdict(int) - mgr._abort_state_lock = threading.Lock() - mgr._deferred_ack_targets = {} - mgr._active_transfer_generations = {room: TRANSFER_GENERATION} - mgr._abort_ack_receipts = set() mgr.is_mla_backend = False mgr.is_hybrid_mla_backend = False mgr.attn_tp_size = 1 mgr.transfer_source_rank = 0 mgr.kv_args = SimpleNamespace(engine_rank=0, kv_data_ptrs=[0]) - mgr._use_torch_transfer = False mgr.exceptions = {} mgr.failure_lock = threading.Lock() mgr.failure_records = {} @@ -735,12 +518,9 @@ def test_given_last_chunk_aborts_mid_transfer_when_worker_finishes_then_failed_s self.assertNotIn(room, mgr.transfer_infos) self.assertNotIn(room, mgr.req_to_decode_prefix_len) mgr.send_aux.assert_called_once() - self.assertEqual( - mgr.send_aux.call_args.args[-1], - completion_notification(21, "aux_nokv_0_0"), - ) + self.assertEqual(mgr.send_aux.call_args.args[-1], "21_aux_nokv_0_0") - def test_given_non_last_chunk_aborts_mid_transfer_then_failed_room_is_retired( + def test_given_non_last_chunk_aborts_mid_transfer_when_worker_finishes_then_failed_status_is_preserved( self, ): room = 22 @@ -751,247 +531,12 @@ def test_given_non_last_chunk_aborts_mid_transfer_then_failed_room_is_retired( self._run_worker_once(mgr, chunk) self.assertEqual(mgr.request_status[room], KVPoll.Failed) - self.assertNotIn(room, mgr.transfer_infos) - self.assertNotIn(room, mgr.req_to_decode_prefix_len) - self.assertIn((room, TRANSFER_GENERATION), mgr._abort_ack_receipts) + self.assertIn(room, mgr.transfer_infos) + self.assertIn(room, mgr.req_to_decode_prefix_len) mgr.send_kvcache.assert_called_once() - def test_waits_for_early_send_event_once_before_source_access(self): - for room, use_torch_transfer in ((23, False), (24, True)): - with self.subTest(use_torch_transfer=use_torch_transfer): - mgr = self._make_manager(room) - mgr._use_torch_transfer = use_torch_transfer - mgr.agent.progress = MagicMock() - actions = [] - captured_handles = [] - mgr.agent.begin_handle_batch = lambda handles: captured_handles.append( - handles - ) - mgr.agent.end_handle_batch = lambda handles: self.assertIs( - captured_handles[-1], handles - ) - wait_event = MagicMock() - wait_event.synchronize.side_effect = lambda: actions.append("wait") - - def send_kvcache(*args, **kwargs): - actions.append("send") - if use_torch_transfer: - captured_handles[-1].append("handle") - return "handle" - - mgr.send_kvcache = MagicMock(side_effect=send_kvcache) - chunk = self._make_chunk(room, [1], is_last_chunk=False) - chunk.wait_event = wait_event - - self._run_worker_once(mgr, chunk) - - self.assertEqual(actions, ["wait", "send"]) - wait_event.synchronize.assert_called_once_with() - self.assertIsNone(chunk.wait_event) - if use_torch_transfer: - mgr.agent.progress.assert_called_once_with() - else: - mgr.agent.progress.assert_not_called() - - def test_torch_failure_is_published_only_after_every_handle_is_drained(self): - room = 26 - mgr = self._make_manager(room) - mgr._use_torch_transfer = True - second_req = replace(mgr.transfer_infos[room]["agent"], agent_name="agent-2") - mgr.transfer_infos[room]["agent-2"] = second_req - mgr.decode_kv_args_table["agent-2"] = mgr.decode_kv_args_table["agent"] - release_kv_cache = MagicMock() - events = [] - batch = [] - original_update_status = mgr.update_status - - def send_kvcache(*args, **kwargs): - handle = "failed" if not batch[-1] else "active" - batch[-1].append(handle) - return handle - - def update_status(candidate_room, status): - events.append(("status", status)) - original_update_status(candidate_room, status) - if status == KVPoll.Failed: - release_kv_cache(candidate_room) - - def cancel_handles_and_wait(handles): - events.append(("drain", tuple(handles))) - self.assertEqual(mgr.request_status[room], KVPoll.Transferring) - release_kv_cache.assert_not_called() - - mgr.update_status = update_status - mgr.send_kvcache = MagicMock(side_effect=send_kvcache) - mgr.agent = SimpleNamespace( - begin_handle_batch=lambda handles: batch.append(handles), - end_handle_batch=lambda handles: self.assertIs(batch[-1], handles), - progress=MagicMock(), - check_xfer_state=MagicMock(return_value="ERR"), - pop_xfer_error=MagicMock(return_value=None), - cancel_handles_and_wait=cancel_handles_and_wait, - ) - chunk = self._make_chunk(room, [1], is_last_chunk=False) - - self._run_worker_once(mgr, chunk) - - self.assertEqual(events[-2], ("drain", ("failed", "active"))) - self.assertEqual(events[-1], ("status", KVPoll.Failed)) - release_kv_cache.assert_called_once_with(room) - self.assertEqual(mgr._staging_outstanding[room], 0) - self.assertFalse(chunk.staging_counted) - - def test_fatal_worker_cut_waits_without_timeout_until_batch_is_quiescent(self): - class FatalWorkerCut(BaseException): - pass - - room = 28 - mgr = self._make_manager(room) - mgr._use_torch_transfer = True - mgr.enable_deferred_decode_kv_release = True - events = [] - batches = [] - cleanup_attempts = 0 - original_update_status = mgr.update_status - - def begin(handles): - batches.append(handles) - events.append("begin") - - def send_then_die(*args, **kwargs): - batches[-1].append("live-handle") - raise FatalWorkerCut("worker cancelled") - - def retain(handles): - self.assertIs(handles, batches[-1]) - self.assertEqual(mgr.request_status[room], KVPoll.Transferring) - events.append("retain") - - def drain(handles): - nonlocal cleanup_attempts - cleanup_attempts += 1 - events.append(f"drain-{cleanup_attempts}") - self.assertIs(handles, batches[-1]) - self.assertEqual(mgr.request_status[room], KVPoll.Transferring) - if cleanup_attempts == 1: - raise KeyboardInterrupt("cleanup interrupted") - - def update_status(candidate_room, status): - original_update_status(candidate_room, status) - if status == KVPoll.Failed: - events.append("failed") - - def ack(candidate_room, generation=None): - self.assertEqual(candidate_room, room) - self.assertEqual(generation, TRANSFER_GENERATION) - self.assertEqual(mgr.request_status[room], KVPoll.Failed) - self.assertEqual(mgr._staging_outstanding[room], 0) - self.assertFalse(chunk.staging_counted) - events.append("ack") - - def end(handles): - self.assertIs(handles, batches[-1]) - events.append("end") - - mgr.update_status = update_status - mgr.send_kvcache = MagicMock(side_effect=send_then_die) - mgr._maybe_ack_drained_abort = ack - mgr.agent = SimpleNamespace( - begin_handle_batch=begin, - end_handle_batch=end, - retain_handle_batch=retain, - cancel_handles_and_wait=drain, - handle_batch_is_quiescent=lambda handles: False, - ) - chunk = self._make_chunk(room, [1], is_last_chunk=False) - queue = SimpleNamespace(get=MagicMock(return_value=chunk)) - - with self.assertRaisesRegex(FatalWorkerCut, "worker cancelled"): - mgr.transfer_worker(queue) - - self.assertEqual( - events, - ["begin", "retain", "drain-1", "drain-2", "failed", "ack", "end"], - ) - self.assertEqual(mgr._staging_outstanding[room], 0) - self.assertFalse(chunk.staging_counted) - self.assertIsInstance(mgr.exceptions[room], RuntimeError) - - def test_torch_peer_bootstrap_base_exception_rolls_back_exact_state(self): - class BootstrapCut(BaseException): - pass - - class StoreThenInterrupt(dict): - def __setitem__(self, key, value): - super().__setitem__(key, value) - raise BootstrapCut("manager commit cut") - - room = 29 - mgr = self._make_manager(room) - mgr._use_torch_transfer = True - mgr.attn_tp_size = 1 - mgr.disaggregation_mode = DisaggregationMode.PREFILL - mgr.requires_dcp_relayout = MagicMock(return_value=False) - mgr.prep_handles = {"": "existing-local"} - mgr.decode_kv_args_table = StoreThenInterrupt() - released = [] - rolled_back = [] - mgr.agent = SimpleNamespace( - add_remote_agent=MagicMock(return_value="agent"), - release_dlist_handle=lambda handle: released.append(handle), - rollback_remote_agent=lambda name, metadata: rolled_back.append( - (name, metadata) - ), - ) - mgr._prepare_payload_xfer = lambda peer: mgr.prep_handles.__setitem__( - "agent", "new-remote" - ) - peer = SimpleNamespace( - agent_name="agent", - agent_metadata=b"peer-metadata", - decode_tp_size=1, - dst_dcp_size=1, - dst_dcp_rank=0, - dst_kv_mem_kinds=["VRAM"], - ) - - with self.assertRaisesRegex(BootstrapCut, "manager commit cut"): - mgr._add_remote_peer(peer) - - self.assertNotIn("agent", mgr.decode_kv_args_table) - self.assertEqual(mgr.prep_handles, {"": "existing-local"}) - self.assertEqual(released, ["new-remote"]) - self.assertEqual(rolled_back, [("agent", b"peer-metadata")]) - - def test_partial_multi_submit_uses_worker_owned_handle_sink(self): - room = 27 - mgr = self._make_manager(room) - mgr._use_torch_transfer = True - mgr.decode_kv_args_table["agent"].kv_xfer_segments = [object(), object()] - drained = [] - - def submit_then_fail(_peer, _src, _dst, _notif, handles): - handles.append("first-active-part") - raise RuntimeError("second part failed before publication") - - mgr.send_kvcache_mixed = MagicMock(side_effect=submit_then_fail) - mgr.agent = SimpleNamespace( - begin_handle_batch=lambda handles: None, - end_handle_batch=lambda handles: None, - progress=MagicMock(), - check_xfer_state=MagicMock(), - pop_xfer_error=MagicMock(return_value=None), - cancel_handles_and_wait=lambda handles: drained.append(tuple(handles)), - ) - chunk = self._make_chunk(room, [1], is_last_chunk=False) - - self._run_worker_once(mgr, chunk) - - self.assertEqual(drained, [("first-active-part",)]) - self.assertEqual(mgr.request_status[room], KVPoll.Failed) - def test_dcp_destinations_use_disjoint_pack_regions_before_chunk_barrier(self): - room = 25 + room = 23 mgr = self._make_manager(room) agents = ("agent0a", "agent0b", "agent1") dcp_ranks = (0, 0, 1) @@ -1005,7 +550,6 @@ def test_dcp_destinations_use_disjoint_pack_regions_before_chunk_barrier(self): dst_aux_index=0, required_dst_info_num=len(agents), dst_state_indices=[], - transfer_generation=TRANSFER_GENERATION, ) for i, agent in enumerate(agents) } @@ -1079,23 +623,10 @@ def check_xfer_state(_handle): class TestNixlNotifications(CustomTestCase): - def _make_manager( - self, - messages, - *, - room=5, - required=None, - generation=TRANSFER_GENERATION, - expected_ranks=frozenset({0}), - ): + def _make_manager(self, messages, required=None): mgr = object.__new__(NixlKVManager) mgr.agent = NotificationFakeAgent(messages) - mgr.transfer_statuses = { - room: TransferStatus( - transfer_generation=generation, - expected_source_ranks=expected_ranks, - ) - } + mgr.transfer_statuses = defaultdict(TransferStatus) mgr.required_prefill_response_num_table = required or {} mgr.enable_staging = False mgr._staging_handler = None @@ -1103,7 +634,7 @@ def _make_manager( return mgr def test_kv_last_notification_sets_expected_count(self): - mgr = self._make_manager([completion_notification(5, "kv_2_1_0")]) + mgr = self._make_manager(["5_kv_2_1_0"]) mgr.update_transfer_status() @@ -1113,26 +644,19 @@ def test_kv_last_notification_sets_expected_count(self): self.assertEqual(status.num_pp_ranks_expected, 1) def test_staging_notification_preserves_agent_name_with_underscores(self): - mgr = self._make_manager( - [completion_notification(5, "stg_0_1_0_2_4_8_agent_with_underscores")] - ) + mgr = self._make_manager(["5_stg_0_1_0_2_4_8_agent_with_underscores"]) calls = [] mgr._handle_staging_chunk_arrived = lambda *args: calls.append(args) mgr.update_transfer_status() - self.assertEqual(calls, [(5, 2, 4, 8, 0)]) + self.assertEqual(calls, [(5, 2, 4, 8, "agent_with_underscores")]) status = mgr.transfer_statuses[5] self.assertEqual(status.received_kvs_per_pp[0], {0}) self.assertEqual(status.expected_kvs_per_pp[0], 1) def test_aux_nokv_marks_zero_expected_chunks_for_pp_rank(self): - mgr = self._make_manager( - [completion_notification(6, "aux_nokv_3_0")], - room=6, - required={6: 4}, - expected_ranks=frozenset({0, 1, 2, 3}), - ) + mgr = self._make_manager(["6_aux_nokv_3"], required={6: 4}) mgr.update_transfer_status() @@ -1142,268 +666,21 @@ def test_aux_nokv_marks_zero_expected_chunks_for_pp_rank(self): self.assertEqual(status.num_pp_ranks_expected, 4) def test_state_notification_marks_pp_rank(self): - mgr = self._make_manager( - [completion_notification(7, "state_2_4")], - room=7, - expected_ranks=frozenset({2}), - ) + mgr = self._make_manager(["7_state_2"]) mgr.update_transfer_status() self.assertEqual(mgr.transfer_statuses[7].received_state_per_pp, {2}) def test_aux_nokv_allows_full_hit_completion(self): - mgr = self._make_manager( - [completion_notification(8, "aux_nokv_0_0")], - room=8, - required={8: 1}, - ) + mgr = self._make_manager(["8_aux_nokv_0"], required={8: 1}) mgr.update_transfer_status() self.assertTrue(mgr.transfer_statuses[8].is_done()) - def test_mixed_parts_complete_only_after_all_generation_bound_parts(self): - mgr = self._make_manager( - [ - completion_notification(9, "kv_0_1_0_part_0_2"), - completion_notification(9, "kv_0_1_0_part_1_2"), - ], - room=9, - ) - - mgr.update_transfer_status() - - status = mgr.transfer_statuses[9] - self.assertEqual(status.received_kvs_per_pp[0], {0}) - self.assertEqual(status.expected_kvs_per_pp[0], 1) - - def test_arbitrary_and_duplicate_source_ranks_cannot_satisfy_completion(self): - mgr = self._make_manager( - [ - completion_notification(10, "kv_0_1_99"), - completion_notification(10, "aux_99"), - completion_notification(10, "kv_0_1_0"), - completion_notification(10, "kv_0_1_0"), - completion_notification(10, "aux_0"), - ], - room=10, - expected_ranks=frozenset({0, 1}), - ) - - mgr.update_transfer_status() - - status = mgr.transfer_statuses[10] - self.assertEqual(status.received_kvs_per_pp, {0: {0}}) - self.assertEqual(status.expected_kvs_per_pp, {0: 1}) - self.assertTrue(status.received_aux) - self.assertFalse(status.is_done()) - - def test_old_generation_cannot_mutate_any_completion_family_after_reuse(self): - mgr = self._make_manager( - [ - completion_notification(5, "kv_0_1_0", OLD_TRANSFER_GENERATION), - completion_notification( - 5, "kv_0_1_0_part_0_1", OLD_TRANSFER_GENERATION - ), - completion_notification(5, "state_0_0", OLD_TRANSFER_GENERATION), - completion_notification(5, "aux_nokv_0_0", OLD_TRANSFER_GENERATION), - completion_notification( - 5, - "stg_0_1_0_0_0_1_old_agent", - OLD_TRANSFER_GENERATION, - ), - ] - ) - staging_calls = [] - mgr._handle_staging_chunk_arrived = lambda *args: staging_calls.append(args) - - mgr.update_transfer_status() - - status = mgr.transfer_statuses[5] - self.assertEqual(status.received_kvs_per_pp, {}) - self.assertEqual(status.expected_kvs_per_pp, {}) - self.assertEqual(status.received_state_per_pp, set()) - self.assertFalse(status.received_aux) - self.assertIsNone(status.received_kv_parts_per_pp) - self.assertEqual(staging_calls, []) - - def test_legacy_and_malformed_notifications_never_create_or_mutate_status(self): - mgr = self._make_manager( - [ - "5_kv_0_1_0", - f"nixlv1_{TRANSFER_GENERATION}_5_aux", - completion_notification(5, "kv_-1_1_0"), - completion_notification(5, "kv_0_2_0"), - completion_notification(5, "aux_extra"), - completion_notification(999, "aux_0"), - ] - ) - - mgr.update_transfer_status() - - self.assertEqual(set(mgr.transfer_statuses), {5}) - status = mgr.transfer_statuses[5] - self.assertEqual(status.received_kvs_per_pp, {}) - self.assertEqual(status.expected_kvs_per_pp, {}) - self.assertFalse(status.received_aux) - - -class TestGenerationBoundStagingControl(CustomTestCase): - @staticmethod - def _request(generation, *, legacy=False): - if legacy: - return [b"STAGING_REQ", b"7", b"0", b"1", b"peer", b"3"] - return [ - STAGING_REQ_WIRE_TAG, - generation.encode("ascii"), - b"7", - b"0", - b"1", - b"peer", - b"3", - ] - - def test_stale_and_legacy_staging_requests_cannot_allocate_reused_room(self): - sock = MagicMock() - receiver = SimpleNamespace( - transfer_generation=TRANSFER_GENERATION, - chunk_staging_infos=[], - _connect_to_bootstrap_server=MagicMock( - return_value=(sock, threading.Lock()) - ), - ) - allocator = SimpleNamespace( - assign=MagicMock(return_value=(3, 128, 0)), total_size=1 << 20 - ) - kwargs = dict( - staging_allocator=allocator, - kv_args=SimpleNamespace( - page_size=64, - kv_item_lens=[4096, 4096], - total_kv_head_num=4, - engine_rank=0, - ), - attn_tp_size=16, - prefill_attn_tp_size=1, - kv_buffer_tensors=None, - room_receivers={7: receiver}, - room_bootstrap={7: [{"pp_rank": 3}]}, - ) - - with patch.dict( - sys.modules, - { - "sglang.srt.disaggregation.common.staging_buffer": _fake_staging_buffer_module() - }, - ): - self.assertFalse( - handle_staging_req(self._request(OLD_TRANSFER_GENERATION), **kwargs) - ) - self.assertFalse( - handle_staging_req( - self._request(TRANSFER_GENERATION, legacy=True), **kwargs - ) - ) - - allocator.assign.assert_not_called() - receiver._connect_to_bootstrap_server.assert_not_called() - self.assertEqual(receiver.chunk_staging_infos, []) - - def test_staging_request_and_response_bind_exact_generation(self): - sock = MagicMock() - receiver = SimpleNamespace( - transfer_generation=TRANSFER_GENERATION, - chunk_staging_infos=[], - _connect_to_bootstrap_server=MagicMock( - return_value=(sock, threading.Lock()) - ), - ) - allocator = SimpleNamespace( - assign=MagicMock(return_value=(3, 128, 0)), total_size=1 << 20 - ) - with patch.dict( - sys.modules, - { - "sglang.srt.disaggregation.common.staging_buffer": _fake_staging_buffer_module() - }, - ): - self.assertTrue( - handle_staging_req( - self._request(TRANSFER_GENERATION), - allocator, - SimpleNamespace( - page_size=64, - kv_item_lens=[4096, 4096], - total_kv_head_num=4, - engine_rank=0, - ), - attn_tp_size=16, - prefill_attn_tp_size=1, - kv_buffer_tensors=None, - room_receivers={7: receiver}, - room_bootstrap={7: [{"pp_rank": 3}]}, - ) - ) - - expected_rsp = [ - STAGING_RSP_WIRE_TAG, - TRANSFER_GENERATION.encode("ascii"), - b"7", - b"0", - b"128", - b"0", - b"640", - b"peer", - ] - sock.send_multipart.assert_called_once_with(expected_rsp) - - tinfo = SimpleNamespace( - transfer_generation=TRANSFER_GENERATION, - staging=None, - ) - transfer_infos = {7: {"peer": tinfo}} - stale_rsp = list(expected_rsp) - stale_rsp[1] = OLD_TRANSFER_GENERATION.encode("ascii") - self.assertFalse(handle_staging_rsp(stale_rsp, transfer_infos)) - self.assertIsNone(tinfo.staging) - self.assertFalse( - handle_staging_rsp( - [b"STAGING_RSP", b"7", b"0", b"128", b"0", b"640", b"peer"], - transfer_infos, - ) - ) - self.assertIsNone(tinfo.staging) - - self.assertTrue(handle_staging_rsp(expected_rsp, transfer_infos)) - self.assertEqual(tinfo.staging.offsets, [128]) - self.assertEqual(tinfo.staging.rounds, [0]) - self.assertEqual(tinfo.staging.ends, [640]) - class TestNixlReceiverPoll(CustomTestCase): - def test_supplied_generation_is_shared_and_read_only_across_decode_ranks(self): - mgr = SimpleNamespace( - transfer_statuses={}, - addr_to_rooms_tracker=defaultdict(set), - update_status=MagicMock(), - ) - - first = NixlKVReceiver( - mgr, "prefill:8998", 11, transfer_generation=TRANSFER_GENERATION - ) - second = NixlKVReceiver( - mgr, "prefill:8998", 11, transfer_generation=TRANSFER_GENERATION - ) - - self.assertEqual(first.transfer_generation, TRANSFER_GENERATION) - self.assertEqual(second.transfer_generation, TRANSFER_GENERATION) - self.assertEqual( - mgr.transfer_statuses[11].transfer_generation, TRANSFER_GENERATION - ) - with self.assertRaises(AttributeError): - first.transfer_generation = OLD_TRANSFER_GENERATION - def _make_receiver(self, status=KVPoll.WaitingForInput): mgr = MagicMock() mgr.waiting_timeout = 5 @@ -1421,7 +698,6 @@ def _make_receiver(self, status=KVPoll.WaitingForInput): receiver.init_time = None receiver.conclude_state = None receiver.abort_notified = False - receiver._transfer_generation = TRANSFER_GENERATION receiver._connection_pool_entries = {} return receiver, mgr @@ -1466,9 +742,7 @@ def test_queued_completion_wins_over_waiting_timeout(self, mock_time): receiver, mgr = self._make_receiver(status=KVPoll.WaitingForInput) receiver.started_transfer = True receiver.init_time = 10.0 - mgr.transfer_statuses = { - 11: TransferStatus(transfer_generation=TRANSFER_GENERATION) - } + mgr.transfer_statuses = {11: TransferStatus()} mgr.check_transfer_done.return_value = True self.assertEqual(receiver.poll(), KVPoll.Success) @@ -1483,7 +757,7 @@ def test_transfer_done_returns_success_and_cleans_room_state(self, mock_time): receiver, mgr = self._make_receiver(status=KVPoll.WaitingForInput) receiver.started_transfer = True receiver.init_time = 10.0 - status = TransferStatus(transfer_generation=TRANSFER_GENERATION) + status = TransferStatus() status.received_aux = True status.num_pp_ranks_expected = 1 status.expected_kvs_per_pp[0] = 0 @@ -1638,7 +912,6 @@ def test_prefetch_staging_reqs_marks_room_when_no_peer_needs_staging(self): dst_aux_index=0, required_dst_info_num=1, dst_state_indices=[], - transfer_generation=TRANSFER_GENERATION, ) } } @@ -1664,11 +937,7 @@ def test_do_staging_transfer_requeues_when_allocation_not_ready(self): prefill_aux_index=None, state_indices=None, ) - req = SimpleNamespace( - room=3, - agent_name="decode_agent", - completion_notification_prefix=f"nixlv2_{TRANSFER_GENERATION}_3", - ) + req = SimpleNamespace(room=3, agent_name="decode_agent") queue = FakeQueue() with patch.dict( @@ -1725,13 +994,7 @@ def test_do_staging_transfer_raises_for_oversized_allocation(self): strategy, kv_chunk, kv_chunk.prefill_kv_indices, - SimpleNamespace( - room=3, - agent_name="decode_agent", - completion_notification_prefix=( - f"nixlv2_{TRANSFER_GENERATION}_3" - ), - ), + SimpleNamespace(room=3, agent_name="decode_agent"), SimpleNamespace(), FakeQueue(), ) @@ -1777,21 +1040,14 @@ def test_do_staging_transfer_builds_staging_notification(self): strategy, kv_chunk, kv_chunk.prefill_kv_indices, - SimpleNamespace( - room=3, - agent_name="decode_agent", - completion_notification_prefix=f"nixlv2_{TRANSFER_GENERATION}_3", - ), + SimpleNamespace(room=3, agent_name="decode_agent"), dst_info, FakeQueue(), ) self.assertEqual(handle, "handle") self.assertFalse(deferred) - self.assertEqual( - calls[0][0][8], - completion_notification(3, "stg_7_1_1_2_4_2_decode_agent"), - ) + self.assertEqual(calls[0][0][8], "3_stg_7_1_1_2_4_2_decode_agent") def test_send_kvcache_staged_uses_one_bulk_vram_write(self): mock_gather = MagicMock() @@ -1820,7 +1076,7 @@ def test_send_kvcache_staged_uses_one_bulk_vram_write(self): dst_tp_rank=0, dst_attn_tp_size=1, dst_kv_item_len=128, - notif=completion_notification(3, "stg_0_1_1_0_0_2_decode_agent"), + notif="3_stg_0_1_1_0_0_2_decode_agent", staging_buffer=FakeStagingBuffer(ptr=0x9000, size=1 << 20), ) @@ -1839,7 +1095,7 @@ def test_send_kvcache_staged_uses_one_bulk_vram_write(self): self.assertEqual(agent.initialize_xfer_calls[0][0], "WRITE") self.assertEqual( agent.initialize_xfer_calls[0][-1], - completion_notification(3, "stg_0_1_1_0_0_2_decode_agent").encode("ascii"), + b"3_stg_0_1_1_0_0_2_decode_agent", ) def test_send_kvcache_staged_falls_back_when_prefill_buffer_too_small(self): @@ -1885,44 +1141,6 @@ def prep_xfer_dlist(self, peer_name, array, mem_kind): return f"handle_{len(self.calls)}" -class TestNixlTorchTransferCompactDlist(CustomTestCase): - def test_equal_tp_dlist_uses_one_strided_region_per_allocation(self): - mgr = object.__new__(NixlKVManager) - mgr._use_torch_transfer = True - mgr.agent = DlistCaptureAgent() - - handle = mgr._prep_equal_tp_dlist( - "decode-agent", - kv_ptrs=[0x10000, 0x20000], - kv_item_lens=[64, 128], - kv_data_lens=[640, 1280], - gpu_id=3, - ) - - self.assertEqual(handle, "handle_1") - peer_name, descriptors, memory_type = mgr.agent.calls[0] - self.assertEqual(peer_name, "decode-agent") - self.assertEqual(memory_type, "VRAM") - np.testing.assert_array_equal( - descriptors, - np.asarray( - [ - [0x10000, 64, 3, 64, 10], - [0x20000, 128, 3, 128, 10], - ], - dtype=np.uint64, - ), - ) - np.testing.assert_array_equal( - repeat_indices_over_layers( - np.asarray([1, 9], dtype=np.int32), - num_layers=2, - layer_length=10, - ), - np.asarray([1, 9, 11, 19], dtype=np.int32), - ) - - class TestNixlHeteroTpReplicatedKV(CustomTestCase): """Regression guard for #31295. 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 c88840de147c..752787665e3b 100644 --- a/test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py +++ b/test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py @@ -7,163 +7,78 @@ after its DONE barrier -- never from the bootstrap thread for an active room. """ -import threading import unittest from unittest.mock import MagicMock from sglang.srt.disaggregation.base.conn import KVPoll from sglang.srt.disaggregation.common.conn import CommonKVManager -from sglang.srt.disaggregation.nixl.conn import GUARD, NixlKVManager +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") -GEN_A = "11" * 16 -GEN_B = "22" * 16 - 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._abort_state_lock = threading.Lock() mgr._deferred_ack_targets = {} - mgr._active_transfer_generations = {} - mgr._abort_ack_receipts = set() mgr._staging_outstanding = {} mgr.request_status = {} mgr._sent = [] # Capture acks instead of opening a socket. - mgr._send_abort_ack = lambda ip, port, room, generation: ( - mgr._sent.append((ip, port, room, generation)) or True - ) + mgr._send_abort_ack = lambda ip, port, room: mgr._sent.append((ip, port, room)) return mgr -def _activate_failed(mgr, room, generation=GEN_A): - mgr.register_transfer_generation(room, generation) - mgr.request_status[room] = KVPoll.Failed - return mgr.register_deferred_ack_target(room, generation, "10.0.0.1", 5000) - - class TestDeferredAckTargets(CustomTestCase): def test_ack_held_until_outstanding_drains(self): mgr = _prefill_mgr() - self.assertTrue(_activate_failed(mgr, 7)) + mgr.register_deferred_ack_target(7, "10.0.0.1", 5000) mgr._staging_outstanding[7] = 1 - mgr._maybe_ack_drained_abort(7, GEN_A) + 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, GEN_A) - self.assertEqual(mgr._sent, [("10.0.0.1", 5000, 7, GEN_A)]) - - def test_one_skipped_writer_cannot_erase_an_active_sibling(self): - mgr = _prefill_mgr() - self.assertTrue(_activate_failed(mgr, 70)) - self.assertEqual(mgr._count_transfer_writer(70), 1) - self.assertEqual(mgr._count_transfer_writer(70), 2) - - self.assertEqual(mgr._uncount_transfer_writer(70), 1) - mgr._maybe_ack_drained_abort(70, GEN_A) - self.assertEqual(mgr._sent, []) - self.assertFalse(mgr.mark_transfer_quiescent(70, GEN_A)) - self.assertEqual(mgr.active_transfer_generation(70), GEN_A) - - self.assertEqual(mgr._uncount_transfer_writer(70), 0) - self.assertTrue(mgr.mark_transfer_quiescent(70, GEN_A)) - self.assertEqual(mgr._sent, [("10.0.0.1", 5000, 70, GEN_A)]) - self.assertIsNone(mgr.active_transfer_generation(70)) + 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_transfer_generation(8, GEN_A) - mgr.request_status[8] = KVPoll.Failed - mgr.register_deferred_ack_target(8, GEN_A, "10.0.0.2", 5001) - mgr._maybe_ack_drained_abort(8, GEN_A) - mgr._maybe_ack_drained_abort(8, GEN_A) + 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, GEN_A), mgr._deferred_ack_targets) + self.assertNotIn(8, mgr._deferred_ack_targets) def test_unregistered_room_is_noop(self): mgr = _prefill_mgr() - mgr._maybe_ack_drained_abort(999, GEN_A) + mgr._maybe_ack_drained_abort(999) self.assertEqual(mgr._sent, []) - def test_normal_send_failure_retains_target_and_receipt_for_retry(self): - mgr = _prefill_mgr() - self.assertTrue(_activate_failed(mgr, 9)) - mgr._send_abort_ack = MagicMock(side_effect=[False, True]) - - mgr._maybe_ack_drained_abort(9, GEN_A) - self.assertIn((9, GEN_A), mgr._deferred_ack_targets) - self.assertIn((9, GEN_A), mgr._abort_ack_receipts) - mgr._maybe_ack_drained_abort(9, GEN_A) - self.assertNotIn((9, GEN_A), mgr._deferred_ack_targets) - self.assertEqual(mgr._send_abort_ack.call_count, 2) - - def test_base_exception_after_acceptance_retains_durable_retry_owner(self): - mgr = _prefill_mgr() - self.assertTrue(_activate_failed(mgr, 10)) - accepted = [] - - def accepted_then_interrupted(ip, port, room, generation): - accepted.append((ip, port, room, generation)) - raise KeyboardInterrupt("lost return") - - mgr._send_abort_ack = accepted_then_interrupted - with self.assertRaises(KeyboardInterrupt): - mgr._maybe_ack_drained_abort(10, GEN_A) - self.assertIn((10, GEN_A), mgr._deferred_ack_targets) - self.assertIn((10, GEN_A), mgr._abort_ack_receipts) - - mgr._send_abort_ack = lambda ip, port, room, generation: ( - accepted.append((ip, port, room, generation)) or True - ) - mgr._maybe_ack_drained_abort(10, GEN_A) - self.assertNotIn((10, GEN_A), mgr._deferred_ack_targets) - self.assertEqual(len(accepted), 2) - - def test_receipt_tombstones_are_not_evicted(self): - mgr = _prefill_mgr() - for room in range(2048): - mgr._remember_abort_receipt_unlocked((room, GEN_A)) - self.assertEqual(len(mgr._abort_ack_receipts), 2048) - self.assertIn((0, GEN_A), mgr._abort_ack_receipts) - 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) - def test_ack_target_is_immutable_for_one_generation(self): - mgr = _prefill_mgr() - mgr.register_transfer_generation(12, GEN_A) - self.assertTrue(mgr.register_deferred_ack_target(12, GEN_A, "10.0.0.1", 5000)) - self.assertFalse(mgr.register_deferred_ack_target(12, GEN_A, "10.0.0.2", 5001)) - target = mgr._deferred_ack_targets[(12, GEN_A)] - self.assertEqual((target.decode_ip, target.decode_port), ("10.0.0.1", 5000)) - 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, generation=GEN_A): + 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"), - generation.encode("ascii"), ] def _mgr(self, enabled=True, room=11, status=KVPoll.WaitingForInput): mgr = _prefill_mgr(NixlKVManager, enabled=enabled) - mgr.register_transfer_generation(room, GEN_A) if status is not None: mgr.request_status[room] = status mgr.record_failure = MagicMock() @@ -179,8 +94,7 @@ def test_in_flight_room_registers_target_and_does_not_ack_yet(self): mgr._staging_outstanding[11] = 1 self.assertTrue(mgr._handle_abort_notification(self._abort_msg())) - target = mgr._deferred_ack_targets[(11, GEN_A)] - self.assertEqual((target.decode_ip, target.decode_port), ("10.0.0.3", 6000)) + 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) @@ -191,7 +105,7 @@ def test_quiescent_active_room_acks_without_waiting_for_a_worker_visit(self): mgr = self._mgr() self.assertTrue(mgr._handle_abort_notification(self._abort_msg())) - self.assertEqual(mgr._sent, [("10.0.0.3", 6000, 11, GEN_A)]) + self.assertEqual(mgr._sent, [("10.0.0.3", 6000, 11)]) self.assertEqual(mgr._deferred_ack_targets, {}) def test_worker_skip_before_registration_still_acks(self): @@ -206,22 +120,21 @@ 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.mark_transfer_quiescent(room, GEN_A) + 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, GEN_A)]) + 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.mark_transfer_quiescent(11, GEN_A) 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, GEN_A)]) + 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): @@ -233,25 +146,16 @@ def test_cleared_room_with_outstanding_chunk_does_not_ack(self): self.assertTrue(mgr._handle_abort_notification(self._abort_msg())) self.assertEqual(mgr._sent, []) - self.assertIn((11, GEN_A), mgr._deferred_ack_targets) + self.assertEqual(mgr._deferred_ack_targets, {}) - def test_explicit_feature_false_cannot_disable_safe_ack_protocol(self): + 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, [("10.0.0.3", 6000, 11, GEN_A)]) - self.assertEqual(mgr.request_status[11], KVPoll.Failed) - - def test_old_generation_abort_cannot_mutate_reused_room(self): - mgr = self._mgr() - self.assertTrue( - mgr._handle_abort_notification(self._abort_msg(generation=GEN_B)) - ) - self.assertEqual(mgr.request_status[11], KVPoll.WaitingForInput) - self.assertEqual(mgr.active_transfer_generation(11), GEN_A) 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. @@ -267,139 +171,17 @@ def test_non_abort_message_is_not_claimed(self): class TestNixlDecodeAckIngest(CustomTestCase): def test_abort_ack_is_aggregated_per_rank(self): - mgr = _prefill_mgr() + # Mirrors the decode listener thread's ABORT_ACK branch. + mgr = CommonKVManager.__new__(CommonKVManager) mgr._deferred_abort_ack_tracker = {} - mgr.register_deferred_abort_room(21, GEN_A, frozenset({0, 1})) - - for rank in (0, 1, 1): - mgr.note_abort_ack(21, GEN_A, rank) + mgr.register_deferred_abort_room(21) - self.assertTrue(mgr.is_abort_release_safe(21, GEN_A)) - - def test_late_metadata_does_not_poison_new_generation(self): - mgr = _prefill_mgr(NixlKVManager) - mgr._bootstrap_errors = {} - mgr.exceptions = {} - mgr.transfer_infos = {21: {}} - mgr.record_failure = MagicMock() - mgr.register_transfer_generation(21, GEN_B) - late_payload = [ - b"21", - b"endpoint", - b"1234", - b"agent", - b"", - b"0", - b"1", - b"", - b"", - b"0", - GEN_A.encode("ascii"), - ] - - mgr._handle_bootstrap_message([GUARD, *late_payload]) - - self.assertEqual(mgr.active_transfer_generation(21), GEN_B) - self.assertEqual(mgr.transfer_infos[21], {}) - self.assertNotIn(21, mgr.exceptions) - mgr.record_failure.assert_not_called() - - def test_malformed_or_generationless_metadata_cannot_poison_reused_room(self): - for generation_frame in (b"not-a-generation", None): - mgr = _prefill_mgr(NixlKVManager) - mgr._bootstrap_errors = {} - mgr.exceptions = {} - mgr.transfer_infos = {23: {}} - mgr.record_failure = MagicMock() - mgr.register_transfer_generation(23, GEN_B) - payload = [ - b"23", - b"endpoint", - b"1234", - b"agent", - b"", - b"0", - b"1", - b"", - b"", - b"0", - ] - if generation_frame is not None: - payload.append(generation_frame) - - mgr._handle_bootstrap_message([GUARD, *payload]) - - self.assertEqual(mgr.active_transfer_generation(23), GEN_B) - self.assertEqual(mgr.transfer_infos[23], {}) - self.assertNotIn(23, mgr.exceptions) - self.assertNotIn("agent", mgr._bootstrap_errors) - mgr.record_failure.assert_not_called() - - def test_retired_metadata_cannot_resurrect_room_before_reuse(self): - mgr = _prefill_mgr(NixlKVManager) - mgr._bootstrap_errors = {} - mgr.exceptions = {} - mgr.transfer_infos = {} - mgr.record_failure = MagicMock() - mgr.register_transfer_generation(22, GEN_A) - mgr.mark_transfer_quiescent(22, GEN_A) - late_payload = [ - b"22", - b"endpoint", - b"1234", - b"agent", - b"", - b"0", - b"1", - b"", - b"", - b"0", - GEN_A.encode("ascii"), - ] + 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())) - mgr._handle_bootstrap_message([GUARD, *late_payload]) - - self.assertIsNone(mgr.active_transfer_generation(22)) - self.assertNotIn(22, mgr.transfer_infos) - self.assertIn((22, GEN_A), mgr._abort_ack_receipts) - self.assertNotIn(22, mgr.exceptions) - mgr.record_failure.assert_not_called() - - def test_retirement_and_duplicate_registration_are_atomic(self): - mgr = _prefill_mgr(NixlKVManager) - rooms = tuple(range(1000, 1256)) - for room in rooms: - self.assertTrue(mgr.register_transfer_generation(room, GEN_A)) - - operation_start = threading.Barrier(2) - operation_done = threading.Barrier(2) - replay_results = [] - - def retire_generations(): - for room in rooms: - operation_start.wait() - mgr.mark_transfer_quiescent(room, GEN_A) - operation_done.wait() - - def replay_metadata_registration(): - for room in rooms: - operation_start.wait() - replay_results.append(mgr.register_transfer_generation(room, GEN_A)) - operation_done.wait() - - retire_thread = threading.Thread(target=retire_generations) - replay_thread = threading.Thread(target=replay_metadata_registration) - retire_thread.start() - replay_thread.start() - retire_thread.join(timeout=5) - replay_thread.join(timeout=5) - - self.assertFalse(retire_thread.is_alive()) - self.assertFalse(replay_thread.is_alive()) - self.assertEqual(len(replay_results), len(rooms)) - for room in rooms: - self.assertIsNone(mgr.active_transfer_generation(room)) - self.assertIn((room, GEN_A), mgr._abort_ack_receipts) + self.assertFalse(mgr.is_abort_release_safe(21, required_acks=3)) + self.assertTrue(mgr.is_abort_release_safe(21, required_acks=2)) 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 54cdd35ee2e7..c3078d5a3614 100644 --- a/test/registered/unit/disaggregation/test_nixl_sender_failure_cleanup.py +++ b/test/registered/unit/disaggregation/test_nixl_sender_failure_cleanup.py @@ -12,7 +12,6 @@ class TestNixlSenderFailureCleanup(unittest.TestCase): def test_failure_exception_cleans_room_state_before_raising(self): room = 7 - generation = "0123456789abcdef0123456789abcdef" expected_exc = RuntimeError("transfer failed") sender = NixlKVSender.__new__(NixlKVSender) sender.bootstrap_room = room @@ -21,10 +20,7 @@ def test_failure_exception_cleans_room_state_before_raising(self): sender._send_error = None staging_ctx = SimpleNamespace( prefetched_rooms={room, 8}, - prefetch_requested={ - (room, generation, 0, "session-a"), - (8, generation, 0, "session-b"), - }, + prefetch_requested={(room, 0, "session-a"), (8, 0, "session-b")}, ) sender.kv_mgr = SimpleNamespace( enable_staging=True, @@ -49,11 +45,9 @@ def test_failure_exception_cleans_room_state_before_raising(self): self.assertNotIn(room, sender.kv_mgr.exceptions) self.assertNotIn(room, sender.kv_mgr.failure_records) self.assertNotIn(room, staging_ctx.prefetched_rooms) - self.assertNotIn( - (room, generation, 0, "session-a"), staging_ctx.prefetch_requested - ) + self.assertNotIn((room, 0, "session-a"), staging_ctx.prefetch_requested) self.assertIn(8, staging_ctx.prefetched_rooms) - self.assertIn((8, generation, 0, "session-b"), staging_ctx.prefetch_requested) + self.assertIn((8, 0, "session-b"), staging_ctx.prefetch_requested) if __name__ == "__main__": diff --git a/test/registered/unit/disaggregation/test_nixl_torch_transfer.py b/test/registered/unit/disaggregation/test_nixl_torch_transfer.py deleted file mode 100644 index 3c63b1d7bb80..000000000000 --- a/test/registered/unit/disaggregation/test_nixl_torch_transfer.py +++ /dev/null @@ -1,2257 +0,0 @@ -"""CPU tests for the experimental PyTorch transfer compatibility boundary.""" - -# Thread targets intentionally retain BaseException for process-control fault -# injection; nested contexts keep the interruption point visually explicit. -# ruff: noqa: BLE001, SIM117 - -from __future__ import annotations - -import ast -import threading -import time -from enum import Enum, auto -from pathlib import Path -from types import SimpleNamespace -from typing import ClassVar -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest -from sglang.srt.disaggregation.nixl.torch_transfer import ( - TorchTransferAgent, - _load_transfer_api, - _register_nixl_backend, - _WorkHandle, -) -from sglang.test.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=1, suite="base-a-test-cpu") - - -class _WorkState(Enum): - PENDING = auto() - RUNNING = auto() - COMPLETED = auto() - SUCCEEDED = COMPLETED - FAILED = auto() - CANCELLED = auto() - - -class _TransferOp(Enum): - READ = "read" - WRITE = "write" - - -class _Resource: - def __init__(self) -> None: - self.closed = False - self.close_calls = 0 - self.close_errors: list[BaseException] = [] - - def close(self) -> None: - self.close_calls += 1 - if self.close_errors: - raise self.close_errors.pop(0) - self.closed = True - - -class _RawSpan: - instances: ClassVar[list[_RawSpan]] = [] - - def __init__(self, **kwargs) -> None: - self.__dict__.update(kwargs) - self.__class__.instances.append(self) - - -class _Registration(_Resource): - def __init__(self, span: _RawSpan, name: str) -> None: - super().__init__() - self.span = span - self.name = name - self.region_calls = [] - - def region(self, offset: int, nbytes: int, **kwargs): - region = ("local", self.name, offset, nbytes, kwargs) - self.region_calls.append(region) - return region - - -class _Peer(_Resource): - def __init__(self, name: str = "remote-endpoint") -> None: - super().__init__() - self.name = name - self.endpoint_id = "remote-id" - self.incarnation = "remote-incarnation" - self.region_calls = [] - - def unsafe_region(self, address: int, nbytes: int, **kwargs): - region = ("remote", address, nbytes, kwargs) - self.region_calls.append(region) - return region - - -class _Plan(_Resource): - def __init__(self, endpoint, kwargs: dict) -> None: - super().__init__() - self.endpoint = endpoint - self.kwargs = kwargs - self.select_indices_calls = [] - self.submit_and_poll_calls = [] - - def select_indices( - self, - operation, - *, - local_indices, - remote_indices, - ): - self.select_indices_calls.append((operation, local_indices, remote_indices)) - raise AssertionError("one-shot SGLang transfers must not create IndexSelection") - - def submit_indices_and_poll( - self, - operation, - *, - local_indices, - remote_indices, - **kwargs, - ): - local_snapshot = np.asarray(local_indices, dtype=np.int32).copy() - remote_snapshot = np.asarray(remote_indices, dtype=np.int32).copy() - call = { - "operation": operation, - "local_indices": tuple(local_snapshot), - "remote_indices": tuple(remote_snapshot), - **kwargs, - } - self.submit_and_poll_calls.append(call) - submission_kwargs = { - key: value - for key, value in kwargs.items() - if key not in {"max_polls", "timeout_ns"} - } - work = self.endpoint._submit( - operation, - self, - { - "local_indices": local_snapshot, - "remote_indices": remote_snapshot, - **submission_kwargs, - }, - ) - endpoint = self.endpoint - state = ( - endpoint.next_fused_state - if endpoint.capabilities.fused_prevalidated_submit_poll - else _WorkState.PENDING - ) - work.done = state in ( - _WorkState.COMPLETED, - _WorkState.FAILED, - _WorkState.CANCELLED, - ) - work.state = state - work.error = endpoint.next_fused_error - endpoint._raise_postcommit_error() - return work, state - - -class _Selection: - def __init__(self, plan, operation, local_indices, remote_indices) -> None: - self.plan = plan - self.operation = operation - self.local_indices = np.asarray(local_indices, dtype=np.int32).copy() - self.remote_indices = np.asarray(remote_indices, dtype=np.int32).copy() - self.submit_and_poll_calls = [] - - def submit_and_poll(self, **kwargs): - call = dict(kwargs) - self.submit_and_poll_calls.append(call) - self.plan.submit_and_poll_calls.append((self, call)) - submission_kwargs = { - key: value - for key, value in kwargs.items() - if key not in {"max_polls", "timeout_ns"} - } - work = self.plan.endpoint._submit( - self.operation, - self.plan, - { - "local_indices": self.local_indices, - "remote_indices": self.remote_indices, - **submission_kwargs, - }, - ) - endpoint = self.plan.endpoint - state = ( - endpoint.next_fused_state - if endpoint.capabilities.fused_prevalidated_submit_poll - else _WorkState.PENDING - ) - work.done = state in ( - _WorkState.COMPLETED, - _WorkState.FAILED, - _WorkState.CANCELLED, - ) - work.state = state - work.error = endpoint.next_fused_error - endpoint._raise_postcommit_error() - return work, state - - -class _Work(_Resource): - def __init__(self, recovery_token=None) -> None: - super().__init__() - self.recovery_token = recovery_token - self.done = False - self._state = _WorkState.PENDING - self.error = None - self.cancel_completes = True - self.test_error: Exception | None = None - self.test_calls = 0 - self.state_hook = None - - @property - def state(self) -> _WorkState: - self.test_calls += 1 - if self.state_hook is not None: - self.state_hook() - if self.test_error is not None: - raise self.test_error - return self._state - - @state.setter - def state(self, value: _WorkState) -> None: - self._state = value - - def test(self) -> bool: - self.test_calls += 1 - if self.test_error is not None: - raise self.test_error - return self.done - - def cancel(self) -> bool: - if self.cancel_completes: - self.done = True - self.state = _WorkState.CANCELLED - return self.cancel_completes - - def close(self) -> None: - super().close() - if self.closed: - self.recovery_token = None - - -class _ProgressMode(Enum): - MANUAL = auto() - BACKGROUND = auto() - - -class _ThreadMode(Enum): - CALLER_SERIALIZED = auto() - SERIALIZED = auto() - MULTIPLE = auto() - - -class _Endpoint(_Resource): - instances: ClassVar[list[_Endpoint]] = [] - synchronous_notification_send = True - constructor_cut: tuple[BaseException, bool, list[BaseException]] | None = None - missing_capability: str | None = None - - def __init__( - self, - name: str, - *, - backend: str, - endpoint_id: str, - progress_mode: _ProgressMode, - thread_mode: _ThreadMode, - options: dict, - ) -> None: - super().__init__() - self.name = name - self.backend = backend - self.endpoint_id = endpoint_id - self.progress_mode = progress_mode - self.thread_mode = thread_mode - self.options = options - self.capabilities = SimpleNamespace( - notifications=True, - attached_notifications=True, - standalone_notifications=True, - raw_spans=True, - strided_regions=True, - write=True, - fused_prevalidated_submit_poll=True, - synchronous_notification_send=self.__class__.synchronous_notification_send, - ) - self.registrations = [] - self.peers = [] - self.plans = [] - self.bound_transfers = [] - self.transfer_slots = [] - self.submissions = [] - self.notifications = [] - self.synchronous_notifications = [] - self.incoming_notifications = [] - self.progress_calls = 0 - self.next_fused_result = False - self.next_fused_state = _WorkState.PENDING - self.next_fused_error = None - self.poison_index_submit = False - self.submit_attempts = 0 - self.submit_error: BaseException | None = None - self.postcommit_error: BaseException | None = None - self.postcommit_hook = None - self.open_works_errors: list[BaseException] = [] - self.next_work_done = False - self.next_work_state = _WorkState.PENDING - self._open_works = [] - self.open_works_reads = 0 - self.construction_snapshot_reads = 0 - self.poison_construction_snapshots = False - missing_capability = self.__class__.missing_capability - if missing_capability is not None: - setattr(self.capabilities, missing_capability, False) - self.__class__.instances.append(self) - constructor_cut = self.__class__.constructor_cut - if constructor_cut is not None: - self.__class__.constructor_cut = None - error, mismatch, close_errors = constructor_cut - self.close_errors.extend(close_errors) - if mismatch: - self.endpoint_id = f"foreign-{endpoint_id}" - error.recovery_endpoint = self - raise error - - @property - def open_registrations(self): - self._record_construction_snapshot() - return tuple(item for item in self.registrations if not item.closed) - - @property - def open_peers(self): - self._record_construction_snapshot() - return tuple(item for item in self.peers if not item.closed) - - @property - def open_plans(self): - self._record_construction_snapshot() - return tuple(item for item in self.plans if not item.closed) - - @property - def open_bound_transfers(self): - self._record_construction_snapshot() - return tuple(item for item in self.bound_transfers if not item.closed) - - @property - def open_transfer_slots(self): - self._record_construction_snapshot() - return tuple(item for item in self.transfer_slots if not item.closed) - - def _record_construction_snapshot(self): - self.construction_snapshot_reads += 1 - if self.poison_construction_snapshots: - raise AssertionError("clean path read a Core construction snapshot") - - @property - def open_works(self): - self.open_works_reads += 1 - if self.open_works_errors: - raise self.open_works_errors.pop(0) - return tuple(work for work in self._open_works if not work.closed) - - def register(self, span, *, name): - registration = _Registration(span, name) - self.registrations.append(registration) - return registration - - def export_metadata(self, registrations): - assert registrations == list(self.open_registrations) - return b"endpoint-metadata" - - def import_peer(self, metadata, **kwargs): - assert metadata == b"peer-metadata" - if kwargs: - assert kwargs == {"expected_name": "decode-agent"} - peer = _Peer() - self.peers.append(peer) - return peer - - def prepare(self, **kwargs): - plan = _Plan(self, kwargs) - self.plans.append(plan) - return plan - - def _submit(self, operation, plan, kwargs): - self.submit_attempts += 1 - if self.submit_error is not None: - raise self.submit_error - work = _Work(kwargs.get("recovery_token")) - work.done = self.next_work_done - work.state = self.next_work_state - self.submissions.append((operation, plan, kwargs, work)) - self._open_works.append(work) - return work - - def _raise_postcommit_error(self): - if self.postcommit_hook is not None: - self.postcommit_hook() - error, self.postcommit_error = self.postcommit_error, None - if error is not None: - raise error - - def write(self, plan, **kwargs): - work = self._submit("WRITE", plan, kwargs) - self._raise_postcommit_error() - return work - - def read(self, plan, **kwargs): - work = self._submit("READ", plan, kwargs) - self._raise_postcommit_error() - return work - - def progress(self): - self.progress_calls += 1 - - def send_notification(self, peer, payload, *, recovery_token=None): - work = _Work(recovery_token) - self.notifications.append((peer, payload, recovery_token, work)) - self._open_works.append(work) - self._raise_postcommit_error() - return work - - def send_notification_sync(self, peer, payload): - self.synchronous_notifications.append((peer, payload)) - - def poll_notifications(self): - result = self.incoming_notifications - self.incoming_notifications = [] - return result - - -def _make_api(): - return SimpleNamespace( - BACKEND_FACTORY_API_VERSION=2, - BackendFactoryV2=object, - Endpoint=_Endpoint, - ProgressMode=_ProgressMode, - RawSpan=_RawSpan, - ThreadMode=_ThreadMode, - TransferOp=_TransferOp, - TransferPlan=_Plan, - IndexSelection=_Selection, - WorkState=_WorkState, - ) - - -def _reset_fake_api() -> None: - _Endpoint.instances.clear() - _Endpoint.constructor_cut = None - _Endpoint.missing_capability = None - _Endpoint.synchronous_notification_send = True - _RawSpan.instances.clear() - - -def _make_uncreated_agent(): - _reset_fake_api() - with patch( - "sglang.srt.disaggregation.nixl.torch_transfer._load_transfer_api", - return_value=_make_api(), - ): - return TorchTransferAgent( - "local-endpoint", - owner=object(), - background_progress=False, - endpoint_backend="fake", - ) - - -def _make_agent( - *, - background_progress: bool = False, - synchronous_notification_send: bool = True, -): - _reset_fake_api() - _Endpoint.synchronous_notification_send = synchronous_notification_send - api = _make_api() - owner = object() - with ( - patch( - "sglang.srt.disaggregation.nixl.torch_transfer._load_transfer_api", - return_value=api, - ), - patch( - "sglang.srt.disaggregation.nixl.torch_transfer._register_nixl_backend" - ) as register_nixl, - ): - agent = TorchTransferAgent( - "local-endpoint", - owner=owner, - background_progress=background_progress, - num_threads=8 if background_progress else 0, - ) - agent.create_backend("UCX", {"num_threads": "8"}) - register_nixl.assert_called_once_with() - return agent, _Endpoint.instances[-1], owner - - -def _register_and_add_peer(agent): - registrations = agent.register_memory([(0x1000, 0x400, 0, "")], "VRAM") - agent.add_remote_agent(b"peer-metadata", peer_name="decode-agent") - return registrations - - -def test_raw_registration_keeps_owner_and_exports_metadata(): - agent, endpoint, owner = _make_agent(background_progress=True) - - registrations = agent.register_memory( - [(0x1000, 0x100, 2, ""), (0x2000, 0x80, 0, "")], "VRAM" - ) - - assert endpoint.progress_mode is _ProgressMode.BACKGROUND - assert endpoint.thread_mode is _ThreadMode.MULTIPLE - assert endpoint.options == { - "backends": ["UCX"], - "backend_init_params": {"UCX": {"num_threads": "8"}}, - "transfer_backends": ["UCX"], - "num_threads": 8, - } - assert len(registrations) == 2 - assert _RawSpan.instances[0].device == "cuda:2" - assert _RawSpan.instances[0].owner is owner - assert agent.get_agent_metadata() == b"endpoint-metadata" - - -def test_prepared_write_reuses_indexed_plan_and_attaches_notification(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.prep_xfer_dlist("", [(0x1000, 64, 0), (0x1040, 64, 0)], "VRAM") - remote = agent.prep_xfer_dlist( - "decode-agent", [(0x5000, 64, 0), (0x5040, 64, 0)], "VRAM" - ) - - source_indices = np.asarray([0, 1], dtype=np.int32) - destination_indices = np.asarray([1, 0], dtype=np.int32) - notification = bytearray(b"7_kv_0_1_0") - first = agent.make_prepped_xfer( - "WRITE", - local, - source_indices, - remote, - destination_indices, - notification, - ) - second = agent.make_prepped_xfer("WRITE", local, [1], remote, [0]) - - assert len(endpoint.plans) == 1 - assert endpoint.plans[0].kwargs["indexed"] is True - assert endpoint.submissions == [] - assert endpoint.submit_attempts == 0 - assert first.work is None - assert second.work is None - assert endpoint.plans[0].select_indices_calls == [] - assert first.local_indices is source_indices - assert first.remote_indices is destination_indices - assert source_indices.flags.writeable is False - assert destination_indices.flags.writeable is False - - with pytest.raises(ValueError, match="read-only"): - source_indices[:] = [1, 1] - with pytest.raises(ValueError, match="read-only"): - destination_indices[:] = [0, 0] - notification[:] = b"changed_tag" - - assert agent.transfer(first) == "PROC" - first_submission = endpoint.submissions[0][2] - assert tuple(first_submission["local_indices"]) == (0, 1) - assert tuple(first_submission["remote_indices"]) == (1, 0) - assert first_submission["local_indices"].dtype == np.int32 - assert first_submission["remote_indices"].dtype == np.int32 - assert first_submission["notification"] == b"7_kv_0_1_0" - assert endpoint.plans[0].submit_and_poll_calls[0] == { - "operation": _TransferOp.WRITE, - "local_indices": (0, 1), - "remote_indices": (1, 0), - "notification": b"7_kv_0_1_0", - "max_polls": 0, - "timeout_ns": None, - "recovery_token": first, - } - assert first.local_indices is None - assert first.remote_indices is None - assert endpoint.plans[0].select_indices_calls == [] - assert endpoint.open_works_reads == 0 - assert first.work.test_calls == 0 - first.work.done = True - first.work.state = _WorkState.SUCCEEDED - assert agent.check_xfer_state(first) == "DONE" - assert first.work.test_calls == 1 - assert first.work.closed is True - - assert agent.transfer(second) == "PROC" - assert second.work.test_calls == 0 - assert endpoint.submit_attempts == 2 - assert agent.transfer(second) == "PROC" - assert second.work.test_calls == 1 - assert endpoint.submit_attempts == 2 - second.work.done = True - second.work.state = _WorkState.SUCCEEDED - assert agent.check_xfer_state(second) == "DONE" - assert second.work.test_calls == 2 - agent.release_dlist_handle(local) - assert endpoint.plans[0].closed is True - - -@pytest.mark.parametrize( - ("state", "error", "expected"), - [ - (_WorkState.SUCCEEDED, None, "DONE"), - (_WorkState.FAILED, RuntimeError("fused failure"), "ERR"), - ], -) -def test_indexed_fused_submit_consumes_immediate_terminal_without_state_reread( - state, error, expected -): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - endpoint.next_fused_result = True - endpoint.next_fused_state = state - endpoint.next_fused_error = error - endpoint.poison_index_submit = True - local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") - remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") - - handle = agent.make_prepped_xfer("WRITE", local, [0], remote, [0], b"fused") - assert endpoint.submissions == [] - assert handle.work is None - assert agent.transfer(handle) == expected - assert endpoint.plans[-1].select_indices_calls == [] - assert endpoint.plans[-1].submit_and_poll_calls == [ - { - "operation": _TransferOp.WRITE, - "local_indices": (0,), - "remote_indices": (0,), - "notification": b"fused", - "max_polls": 0, - "timeout_ns": None, - "recovery_token": handle, - } - ] - assert handle.local_indices is None - assert handle.remote_indices is None - assert handle.work.test_calls == 0 - if expected == "ERR": - assert handle.work.closed is True - assert id(handle) not in agent._works - assert agent.check_xfer_state(handle) == expected - assert handle.work.test_calls == 0 - assert handle.work.closed is True - - -def test_work_state_is_one_authoritative_read_then_cached(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") - assert handle.work is None - assert endpoint.submissions == [] - endpoint.next_work_done = True - endpoint.next_work_state = _WorkState.SUCCEEDED - - assert agent.transfer(handle) == "DONE" - handle.work.test = MagicMock(side_effect=AssertionError("test() was queried")) - assert handle.work.test_calls == 1 - assert handle.work.test.call_count == 0 - assert agent.check_xfer_state(handle) == "DONE" - assert handle.work.test_calls == 1 - assert handle.work.closed is True - - -def test_indexed_submit_fallback_is_owned_by_core_api(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - endpoint.capabilities.fused_prevalidated_submit_poll = False - local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") - remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") - - handle = agent.make_prepped_xfer("WRITE", local, [0], remote, [0], b"fallback") - assert endpoint.submissions == [] - assert endpoint.plans[-1].submit_and_poll_calls == [] - assert handle.first_state is None - assert agent.transfer(handle) == "PROC" - assert endpoint.plans[-1].select_indices_calls == [] - assert endpoint.plans[-1].submit_and_poll_calls == [ - { - "operation": _TransferOp.WRITE, - "local_indices": (0,), - "remote_indices": (0,), - "notification": b"fallback", - "max_polls": 0, - "timeout_ns": None, - "recovery_token": handle, - } - ] - assert handle.first_state is None - assert handle.work.test_calls == 0 - assert agent.transfer(handle) == "PROC" - assert handle.work.test_calls == 1 - assert endpoint.submit_attempts == 1 - - -def test_plan_close_failure_quarantines_cache_and_peer_until_retry(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") - remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") - handle = agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) - assert agent.transfer(handle) == "PROC" - handle.work.done = True - handle.work.state = _WorkState.SUCCEEDED - assert agent.check_xfer_state(handle) == "DONE" - plan = endpoint.plans[-1] - plan.close_errors.append(RuntimeError("partially committed close")) - - with pytest.raises(RuntimeError, match="partially committed"): - agent.remove_remote_agent("decode-agent") - - assert not agent._plan_cache - assert agent._closing_plans - assert "decode-agent" in agent._retiring_peers - with pytest.raises(RuntimeError, match="teardown is in progress"): - agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) - - agent.remove_remote_agent("decode-agent") - assert not agent._closing_plans - assert "decode-agent" not in agent._retiring_peers - assert plan.closed is True - - -def test_compact_strided_catalog_preserves_flattened_indices(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.prep_xfer_dlist("", [(0x1000, 32, 0, 64, 8)], "VRAM") - remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0, 64, 8)], "VRAM") - - handle = agent.make_prepped_xfer("WRITE", local, [7, 1], remote, [2, 6], b"compact") - - registration = endpoint.registrations[0] - assert registration.region_calls[-1][4] == {"stride": 64, "count": 8} - assert endpoint.peers[0].region_calls[-1][3] == { - "device": "cuda:0", - "memory_type": "VRAM", - "stride": 64, - "count": 8, - } - assert endpoint.plans[-1].kwargs["local"] == local.regions - assert endpoint.plans[-1].kwargs["remote"] == remote.regions - assert endpoint.submissions == [] - assert agent.transfer(handle) == "PROC" - assert tuple(endpoint.submissions[-1][2]["local_indices"]) == (7, 1) - assert tuple(endpoint.submissions[-1][2]["remote_indices"]) == (2, 6) - assert handle.work.test_calls == 0 - handle.work.done = True - handle.work.state = _WorkState.SUCCEEDED - agent.progress() - assert handle.work.test_calls == 0 - assert agent.check_xfer_state(handle) == "DONE" - assert handle.work.test_calls == 1 - assert handle.work.closed is True - - -def test_status_observation_failure_stays_nonterminal(): - agent, _, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") - assert agent.transfer(handle) == "PROC" - handle.work.test_error = RuntimeError("temporary status failure") - - assert agent.check_xfer_state(handle) == "PROC" - assert handle.terminal_state is None - assert handle.work.closed is False - - handle.work.test_error = None - handle.work.done = True - handle.work.state = _WorkState.SUCCEEDED - agent.progress() - assert handle.work.test_calls == 2 - assert agent.check_xfer_state(handle) == "DONE" - assert handle.terminal_state == "DONE" - assert handle.work.closed is True - - -def test_batch_progress_leaves_one_status_probe_per_active_work(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - handles = [ - agent.initialize_xfer("WRITE", local, remote, "decode-agent") for _ in range(3) - ] - assert [agent.transfer(handle) for handle in handles] == ["PROC"] * 3 - handles[0].work.done = True - handles[0].work.state = _WorkState.SUCCEEDED - - agent.progress() - assert [handle.work.test_calls for handle in handles] == [1, 1, 1] - states = [agent.check_xfer_state(handle) for handle in handles] - - assert endpoint.progress_calls == 1 - assert states == ["DONE", "PROC", "PROC"] - assert [handle.work.test_calls for handle in handles] == [2, 2, 2] - assert handles[0].work.closed is True - assert handles[0].owned_plan.closed is True - assert handles[1].work.closed is False - assert handles[1].owned_plan.closed is False - assert handles[2].work.closed is False - assert handles[2].owned_plan.closed is False - assert set(agent._works) == {id(handles[1]), id(handles[2])} - - handles[1].work.done = True - handles[1].work.state = _WorkState.SUCCEEDED - agent.progress() - assert [handle.work.test_calls for handle in handles] == [2, 2, 2] - states = [agent.check_xfer_state(handle) for handle in handles[1:]] - - assert endpoint.progress_calls == 2 - assert states == ["DONE", "PROC"] - assert [handle.work.test_calls for handle in handles] == [2, 3, 3] - assert handles[1].work.closed is True - assert handles[1].owned_plan.closed is True - assert handles[2].work.closed is False - assert handles[2].owned_plan.closed is False - assert set(agent._works) == {id(handles[2])} - - -def test_detached_notification_work_is_still_reaped_by_progress(): - agent, endpoint, _ = _make_agent(synchronous_notification_send=False) - _register_and_add_peer(agent) - handle = agent.send_notif("decode-agent", b"abort") - assert endpoint.notifications[-1][2] is handle - assert endpoint.open_works_reads == 0 - handle.work.done = True - handle.work.state = _WorkState.SUCCEEDED - - agent.progress() - - assert endpoint.progress_calls == 1 - assert handle.work.test_calls == 1 - assert handle.work.closed is True - assert handle.terminal_state == "DONE" - assert agent._works == {} - assert agent._detached_works == [] - - -def test_uncancellable_sibling_is_retained_then_reaped_when_terminal(): - agent, _, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - failed, sibling, foreground = [ - agent.initialize_xfer("WRITE", local, remote, "decode-agent") for _ in range(3) - ] - assert [agent.transfer(handle) for handle in (failed, sibling, foreground)] == [ - "PROC", - "PROC", - "PROC", - ] - failed.work.done = True - failed.work.state = _WorkState.FAILED - sibling.work.cancel_completes = False - - agent.cancel_handles([failed, sibling]) - - assert failed.work.closed is True - assert sibling.work.closed is False - assert id(sibling) in agent._works - assert id(sibling) in agent._retained_work_ids - sweep_thread = agent._sweep_thread - assert sweep_thread is not None - foreground_probe_count = foreground.work.test_calls - - sibling.work.done = True - sibling.work.state = _WorkState.SUCCEEDED - deadline = time.monotonic() + 1 - while not sibling.work.closed and time.monotonic() < deadline: - time.sleep(0.01) - - assert sibling.work.closed is True - sweep_thread.join(timeout=1) - assert sweep_thread.is_alive() is False - assert agent._retained_work_ids == set() - assert foreground.work.test_calls == foreground_probe_count - foreground.work.done = True - foreground.work.state = _WorkState.SUCCEEDED - assert agent.check_xfer_state(foreground) == "DONE" - - -def test_cancel_handles_and_wait_does_not_return_before_uncancellable_work_is_safe(): - agent, _, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - failed, sibling = [ - agent.initialize_xfer("WRITE", local, remote, "decode-agent") for _ in range(2) - ] - assert agent.transfer(failed) == "PROC" - assert agent.transfer(sibling) == "PROC" - failed.work.done = True - failed.work.state = _WorkState.FAILED - sibling.work.cancel_completes = False - - def complete_sibling() -> None: - time.sleep(0.01) - sibling.work.done = True - sibling.work.state = _WorkState.SUCCEEDED - - completion = threading.Thread(target=complete_sibling) - completion.start() - agent.cancel_handles_and_wait([failed, sibling]) - completion.join(timeout=1) - - assert completion.is_alive() is False - assert failed.work.closed is True - assert sibling.work.closed is True - assert agent._works == {} - - -def test_background_progress_does_not_call_manual_progress(): - agent, endpoint, _ = _make_agent(background_progress=True) - _register_and_add_peer(agent) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") - - class _UnexpectedLock: - def __enter__(self): - raise AssertionError("empty background progress acquired the lock") - - def __exit__(self, *args): - return False - - with patch.object(agent, "_lock", _UnexpectedLock()): - agent.progress() - - assert endpoint.progress_calls == 0 - with pytest.raises(RuntimeError, match="has not been posted"): - agent.check_xfer_state(handle) - assert endpoint.submissions == [] - assert agent.transfer(handle) == "PROC" - assert handle.work.test_calls == 1 - - handle.work.done = True - handle.work.state = _WorkState.SUCCEEDED - agent.progress() - - assert endpoint.progress_calls == 0 - assert handle.work.test_calls == 1 - assert agent.check_xfer_state(handle) == "DONE" - assert handle.work.test_calls == 2 - assert handle.work.closed is True - assert handle.owned_plan.closed is True - - -def test_multiple_mode_admission_gate_drains_progress_before_close(): - agent, endpoint, _ = _make_agent(background_progress=False) - progress_entered = threading.Event() - progress_release = threading.Event() - close_started = threading.Event() - close_finished = threading.Event() - failures: list[BaseException] = [] - original_progress = endpoint.progress - - def blocking_progress(): - progress_entered.set() - if not progress_release.wait(1): - raise RuntimeError("test progress release timed out") - original_progress() - - endpoint.progress = blocking_progress - - def run_progress(): - try: - agent.progress() - except BaseException as error: - failures.append(error) - - def run_close(): - close_started.set() - try: - agent.close() - except BaseException as error: - failures.append(error) - finally: - close_finished.set() - - progress_thread = threading.Thread(target=run_progress) - close_thread = threading.Thread(target=run_close) - progress_thread.start() - assert progress_entered.wait(1) - close_thread.start() - assert close_started.wait(1) - assert not close_finished.wait(0.05) - progress_release.set() - progress_thread.join(timeout=1) - close_thread.join(timeout=1) - - assert not progress_thread.is_alive() - assert not close_thread.is_alive() - assert failures == [] - assert endpoint.closed is True - - -def test_distinct_handle_submits_overlap_without_adapter_locks(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") - remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") - handles = [ - agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) for _ in range(2) - ] - barrier = threading.Barrier(2) - counter_lock = threading.Lock() - active = 0 - max_active = 0 - original_submit = endpoint._submit - - def overlapping_submit(operation, plan, kwargs): - nonlocal active, max_active - assert not agent._lock._is_owned() - assert not kwargs["recovery_token"]._state_lock.locked() - with counter_lock: - active += 1 - max_active = max(max_active, active) - try: - barrier.wait(timeout=1) - return original_submit(operation, plan, kwargs) - finally: - with counter_lock: - active -= 1 - - endpoint._submit = overlapping_submit - results = [] - failures = [] - - def submit(handle): - try: - results.append(agent.transfer(handle)) - except BaseException as error: - failures.append(error) - - threads = [threading.Thread(target=submit, args=(handle,)) for handle in handles] - for thread in threads: - thread.start() - for thread in threads: - thread.join(timeout=2) - - assert all(not thread.is_alive() for thread in threads) - assert failures == [] - assert sorted(results) == ["PROC", "PROC"] - assert max_active == 2 - - -def test_distinct_handle_status_refreshes_overlap_without_adapter_locks(): - agent, _, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - handles = [ - agent.initialize_xfer("WRITE", local, remote, "decode-agent") for _ in range(2) - ] - assert [agent.transfer(handle) for handle in handles] == ["PROC", "PROC"] - barrier = threading.Barrier(2) - counter_lock = threading.Lock() - active = 0 - max_active = 0 - - def make_hook(handle): - def hook(): - nonlocal active, max_active - assert not agent._lock._is_owned() - assert not handle._state_lock.locked() - with counter_lock: - active += 1 - max_active = max(max_active, active) - try: - barrier.wait(timeout=1) - finally: - with counter_lock: - active -= 1 - - return hook - - for handle in handles: - handle.work.state_hook = make_hook(handle) - results = [] - failures = [] - - def poll(handle): - try: - results.append(agent.check_xfer_state(handle)) - except BaseException as error: - failures.append(error) - - threads = [threading.Thread(target=poll, args=(handle,)) for handle in handles] - for thread in threads: - thread.start() - for thread in threads: - thread.join(timeout=2) - - assert all(not thread.is_alive() for thread in threads) - assert failures == [] - assert sorted(results) == ["PROC", "PROC"] - assert max_active == 2 - - -def test_teardown_gate_never_refreshes_status_under_agent_lock(): - agent, _, _ = _make_agent() - registrations = _register_and_add_peer(agent) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") - assert agent.transfer(handle) == "PROC" - with handle._state_lock: - handle.submission_interrupted = True - probes_before = handle.work.test_calls - open_works_before = agent._endpoint.open_works_reads - handle.work.state_hook = lambda: pytest.fail( - "teardown gate must not perform a Core status read" - ) - - with pytest.raises(RuntimeError, match="submission.*unsettled"): - agent.deregister_memory(registrations) - - assert handle.work.test_calls == probes_before - assert agent._endpoint.open_works_reads == open_works_before - - -def test_concurrent_interrupted_recovery_has_one_core_snapshot_owner(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - handle = _WorkHandle(phase="recovery_pending") - handle.submission_interrupted = True - agent._track_work(handle) - recovered = _Work(handle) - endpoint._open_works.append(recovered) - entered = threading.Event() - release = threading.Event() - original_open_works = type(endpoint).open_works.fget - - def blocking_open_works(instance): - entered.set() - assert release.wait(1) - return original_open_works(instance) - - results = [] - failures = [] - - def recover(): - try: - results.append(agent._recover_interrupted_post_work(handle)) - except BaseException as error: - failures.append(error) - - with patch.object(type(endpoint), "open_works", property(blocking_open_works)): - owner = threading.Thread(target=recover) - contender = threading.Thread(target=recover) - owner.start() - assert entered.wait(1) - contender.start() - contender.join(timeout=1) - release.set() - owner.join(timeout=1) - - assert not owner.is_alive() - assert not contender.is_alive() - assert failures == [] - assert sorted(results) == ["pending", "recovered"] - assert endpoint.open_works_reads == 1 - assert handle.work is recovered - assert handle.phase == "posted" - - -def test_manual_progress_does_not_hold_adapter_state_lock(): - agent, endpoint, _ = _make_agent(background_progress=False) - observed = [] - - def inspect_progress_lock(): - observed.append(agent._lock._is_owned()) - - endpoint.progress = inspect_progress_lock - agent.progress() - assert observed == [False] - - -def test_close_adopts_authoritative_closed_state_before_child_queries(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - snapshots_before = endpoint.construction_snapshot_reads - endpoint.poison_construction_snapshots = True - endpoint.closed = True - - agent.close() - - assert endpoint.construction_snapshot_reads == snapshots_before - assert agent._endpoint is None - assert agent._works == {} - assert agent._peers == {} - - -def test_close_commit_lost_return_converges_on_retry(): - agent, endpoint, _ = _make_agent() - - def committed_then_interrupted(): - endpoint.close_calls += 1 - endpoint.closed = True - raise KeyboardInterrupt("close return lost") - - endpoint.close = committed_then_interrupted - with pytest.raises(KeyboardInterrupt, match="close return lost"): - agent.close() - - assert agent._endpoint is None - assert endpoint.close_calls == 1 - agent.close() - assert endpoint.close_calls == 1 - - -def test_background_progress_reaps_nonempty_detached_work(): - agent, endpoint, _ = _make_agent( - background_progress=True, synchronous_notification_send=False - ) - _register_and_add_peer(agent) - handle = agent.send_notif("decode-agent", b"abort") - handle.work.done = True - handle.work.state = _WorkState.SUCCEEDED - - agent.progress() - - assert endpoint.progress_calls == 0 - assert handle.work.test_calls == 1 - assert handle.work.closed is True - assert agent._detached_works == [] - - -def test_conn_enables_background_progress_for_prefill_and_decode_roles(): - repo_root = Path(__file__).resolve().parents[4] - source = (repo_root / "python/sglang/srt/disaggregation/nixl/conn.py").read_text() - tree = ast.parse(source) - calls = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id == "TorchTransferAgent" - ] - assert len(calls) == 1 - keywords = {keyword.arg: keyword.value for keyword in calls[0].keywords} - assert isinstance(keywords["background_progress"], ast.Constant) - assert keywords["background_progress"].value is True - assert isinstance(keywords["num_threads"], ast.Name) - assert keywords["num_threads"].id == "num_threads" - - assignments = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.Assign) - and any( - isinstance(target, ast.Name) and target.id == "num_threads" - for target in node.targets - ) - ] - role_assignment = next( - node for node in assignments if isinstance(node.value, ast.IfExp) - ) - assert isinstance(role_assignment.value.body, ast.Constant) - assert role_assignment.value.body.value == 8 - assert isinstance(role_assignment.value.orelse, ast.Constant) - assert role_assignment.value.orelse.value == 0 - - -def test_registration_names_do_not_collide_after_deregister(): - agent, endpoint, _ = _make_agent() - first = agent.register_memory([(0x1000, 0x100, 0, "")], "VRAM") - second = agent.register_memory([(0x2000, 0x100, 0, "")], "VRAM") - agent.deregister_memory(first) - third = agent.register_memory([(0x3000, 0x100, 0, "")], "VRAM") - - assert second[0].name == "sglang_region_1" - assert third[0].name == "sglang_region_2" - assert len({registration.name for registration in endpoint.registrations}) == 3 - assert agent._registration_groups == [second, third] - - -@pytest.mark.parametrize("operation", ["READ", "WRITE"]) -def test_one_off_transfer_orients_local_and_remote_regions(operation): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.get_xfer_descs([(0x1080, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5080, 32, 0)], "VRAM") - source, destination = (remote, local) if operation == "READ" else (local, remote) - - handle = agent.initialize_xfer( - operation, source, destination, "decode-agent", b"done" - ) - - plan = endpoint.plans[-1] - assert plan.kwargs["indexed"] is False - assert plan.kwargs["local"][0][0] == "local" - assert plan.kwargs["remote"][0][0] == "remote" - assert endpoint.submissions == [] - assert agent.transfer(handle) == "PROC" - assert endpoint.submissions[-1][0] == operation - assert endpoint.submissions[-1][2]["notification"] == b"done" - assert endpoint.submissions[-1][2]["recovery_token"] is handle - assert endpoint.open_works_reads == 0 - handle.work.done = True - handle.work.state = _WorkState.SUCCEEDED - assert agent.check_xfer_state(handle) == "DONE" - assert plan.closed is True - - -def test_standalone_notifications_preserve_peer_alias(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - notification = SimpleNamespace( - source_endpoint_id="remote-id", - source_incarnation="remote-incarnation", - payload=b"9_aux", - ) - endpoint.incoming_notifications.append(notification) - - payload = b"abort" - result = agent.send_notif("decode-agent", payload) - - assert agent.get_new_notifs() == {"decode-agent": [b"9_aux"]} - assert endpoint.progress_calls == 0 - assert result is None - assert endpoint.synchronous_notifications[0][1] is payload - assert endpoint.notifications == [] - assert agent._works == {} - assert agent._detached_works == [] - - -def test_notification_without_imported_source_uses_explicit_sentinel(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - endpoint.incoming_notifications.append( - SimpleNamespace( - source_endpoint_id=None, - source_incarnation=None, - payload=b"9_aux", - ) - ) - - assert agent.get_new_notifs() == {"": [b"9_aux"]} - - -def test_notification_from_stale_incarnation_is_rejected(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - endpoint.incoming_notifications.append( - SimpleNamespace( - source_endpoint_id="remote-id", - source_incarnation="old-incarnation", - payload=b"9_aux", - ) - ) - - with pytest.raises(RuntimeError, match="unknown or stale"): - agent.get_new_notifs() - - -def test_notification_batches_group_payloads_and_accept_anonymous_source(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - - class CountingAliases(dict): - lookups = 0 - - def __getitem__(self, key): - self.lookups += 1 - return super().__getitem__(key) - - aliases = CountingAliases(agent._source_aliases) - agent._source_aliases = aliases - endpoint.poll_notification_batches = lambda: [ - SimpleNamespace( - source_endpoint_id="remote-id", - source_incarnation="remote-incarnation", - payloads=(b"first", b"second"), - ), - SimpleNamespace( - source_endpoint_id=None, - source_incarnation=None, - payloads=(b"anonymous",), - ), - ] - - def reject_scalar_poll(): - raise AssertionError("grouped receive must not retry scalar polling") - - endpoint.poll_notifications = reject_scalar_poll - - assert agent.get_new_notifs() == { - "decode-agent": [b"first", b"second"], - "": [b"anonymous"], - } - assert aliases.lookups == 1 - - -def test_notification_batch_from_stale_incarnation_is_rejected(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - endpoint.poll_notification_batches = lambda: [ - SimpleNamespace( - source_endpoint_id="remote-id", - source_incarnation="old-incarnation", - payloads=(b"stale",), - ) - ] - - with pytest.raises(RuntimeError, match="unknown or stale"): - agent.get_new_notifs() - - -def test_notification_batch_failure_is_never_retried_as_scalar(): - agent, endpoint, _ = _make_agent() - batch_calls = 0 - scalar_calls = 0 - - def fail_batch_poll(): - nonlocal batch_calls - batch_calls += 1 - raise RuntimeError("destructive batch failure") - - def count_scalar_poll(): - nonlocal scalar_calls - scalar_calls += 1 - return [] - - endpoint.poll_notification_batches = fail_batch_poll - endpoint.poll_notifications = count_scalar_poll - - with pytest.raises(RuntimeError, match="destructive batch failure"): - agent.get_new_notifs() - assert batch_calls == 1 - assert scalar_calls == 0 - - -def test_active_work_prevents_unsafe_close(): - agent, _, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") - assert agent.transfer(handle) == "PROC" - handle.work.cancel_completes = False - - with pytest.raises(RuntimeError, match="active transfer"): - agent.close() - - handle.work.done = True - handle.work.state = _WorkState.CANCELLED - agent.close() - - -def test_prepared_release_and_agent_close_never_post(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - local_prepared = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") - remote_prepared = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") - indexed = agent.make_prepped_xfer( - "WRITE", local_prepared, [0], remote_prepared, [0], b"never-posted" - ) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - owned = agent.initialize_xfer( - "WRITE", local, remote, "decode-agent", b"also-never-posted" - ) - - agent.release_xfer_handle(owned) - assert owned.owned_plan.closed is True - assert endpoint.submit_attempts == 0 - assert indexed.local_indices == (0,) - assert indexed.remote_indices == (0,) - agent.close() - - assert indexed.phase == "released" - assert indexed.local_indices is None - assert indexed.remote_indices is None - assert endpoint.submit_attempts == 0 - assert all(plan.closed for plan in endpoint.plans) - - -def test_post_failure_is_terminal_and_never_retried(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") - remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") - handle = agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) - endpoint.submit_error = RuntimeError("post failed") - - with pytest.raises(RuntimeError, match="post failed"): - agent.transfer(handle) - assert endpoint.submit_attempts == 1 - assert handle.terminal_state == "ERR" - assert handle.phase == "released" - assert agent.transfer(handle) == "ERR" - assert endpoint.submit_attempts == 1 - - -@pytest.mark.parametrize("indexed", [True, False], ids=["indexed", "one-off"]) -def test_postcommit_base_exception_recovers_exact_core_work(indexed): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - endpoint.postcommit_error = KeyboardInterrupt("submit return interrupted") - - if indexed: - endpoint.next_fused_state = _WorkState.COMPLETED - local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") - remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") - handle = agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) - else: - endpoint.next_work_done = True - endpoint.next_work_state = _WorkState.COMPLETED - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") - - with pytest.raises(KeyboardInterrupt, match="submit return interrupted"): - agent.transfer(handle) - - recovered = endpoint.submissions[-1][3] - assert handle.work is recovered - assert recovered.close_calls == 1 - assert recovered.closed is True - assert endpoint.open_works == () - assert endpoint.open_works_reads == 2 - assert endpoint.submit_attempts == 1 - assert endpoint.submissions[-1][2]["recovery_token"] is handle - assert handle.terminal_state == "ERR" - assert handle.phase == "released" - assert id(handle) not in agent._works - assert agent.transfer(handle) == "ERR" - assert endpoint.submit_attempts == 1 - if indexed: - assert endpoint.plans[-1].closed is False - else: - assert handle.owned_plan.close_calls == 1 - assert handle.owned_plan.closed is True - - -class _IdentityBombHandle(_WorkHandle): - def __hash__(self): - raise AssertionError("recovery token was hashed") - - def __eq__(self, other): - raise AssertionError("recovery token equality was evaluated") - - def __bool__(self): - raise AssertionError("recovery token truthiness was evaluated") - - -class _UnprintableInterrupt(KeyboardInterrupt): - def __str__(self): - raise RuntimeError("exception formatting is hostile") - - -def test_interrupted_post_ignores_distractor_and_uses_identity_only(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - distractor = _Work(object()) - endpoint._open_works.append(distractor) - endpoint.next_work_done = True - endpoint.next_work_state = _WorkState.COMPLETED - endpoint.postcommit_error = KeyboardInterrupt("submit return interrupted") - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - with patch( - "sglang.srt.disaggregation.nixl.torch_transfer._WorkHandle", - _IdentityBombHandle, - ): - handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") - - with pytest.raises(KeyboardInterrupt, match="submit return interrupted"): - agent.transfer(handle) - - recovered = endpoint.submissions[-1][3] - assert endpoint.submissions[-1][2]["recovery_token"] is handle - assert recovered.closed is True - assert distractor.closed is False - assert handle.terminal_state == "ERR" - assert handle.phase == "released" - - -def test_ambiguous_token_recovery_is_retained_and_retryable(): - agent, endpoint, _ = _make_agent() - registrations = _register_and_add_peer(agent) - endpoint.postcommit_error = KeyboardInterrupt("submit return interrupted") - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") - duplicate = _Work(handle) - - def add_duplicate(): - endpoint._open_works.append(duplicate) - - endpoint.postcommit_hook = add_duplicate - with patch.object(agent, "_start_retained_work_sweeper") as start_sweeper: - assert agent.transfer(handle) == "PROC" - - assert handle.phase == "recovery_pending" - assert handle.work is None - assert id(handle) in agent._works - assert id(handle) in agent._retained_work_ids - assert start_sweeper.call_count == 1 - assert agent.transfer(handle) == "PROC" - with pytest.raises(RuntimeError, match="recovery is pending"): - agent.release_xfer_handle(handle) - with pytest.raises(RuntimeError, match="submission.*unsettled"): - agent.deregister_memory(registrations) - with pytest.raises(RuntimeError, match="submission.*unsettled"): - agent.remove_remote_agent("decode-agent") - with pytest.raises(RuntimeError, match="active transfer"): - agent.close() - - duplicate.close() - assert agent.transfer(handle) == "ERR" - assert endpoint.submit_attempts == 1 - assert handle.phase == "released" - assert endpoint.submissions[-1][3].closed is True - agent.remove_remote_agent("decode-agent") - agent.deregister_memory(registrations) - - -def test_interrupted_open_works_lookup_retries_without_resubmission(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - endpoint.postcommit_error = KeyboardInterrupt("submit return interrupted") - endpoint.open_works_errors.append(RuntimeError("snapshot interrupted")) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") - - with patch.object(agent, "_start_retained_work_sweeper"): - assert agent.transfer(handle) == "PROC" - assert handle.phase == "recovery_pending" - assert agent.transfer(handle) == "ERR" - - assert endpoint.submit_attempts == 1 - assert endpoint.open_works_reads == 2 - assert handle.phase == "released" - assert "KeyboardInterrupt: submit return interrupted" in str( - agent.pop_xfer_error(handle) - ) - - -def test_exact_recovery_keeps_teardown_blocked_until_work_is_terminal(): - agent, endpoint, _ = _make_agent() - registrations = _register_and_add_peer(agent) - endpoint.postcommit_error = KeyboardInterrupt("submit return interrupted") - - def make_recovered_work_uncancellable(): - endpoint.submissions[-1][3].cancel_completes = False - - endpoint.postcommit_hook = make_recovered_work_uncancellable - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") - - with patch.object(agent, "_start_retained_work_sweeper"): - assert agent.transfer(handle) == "PROC" - assert handle.phase == "posted" - assert handle.work.closed is False - with pytest.raises(RuntimeError, match="submission.*unsettled"): - agent.deregister_memory(registrations) - with pytest.raises(RuntimeError, match="submission.*unsettled"): - agent.remove_remote_agent("decode-agent") - - handle.work.done = True - handle.work.state = _WorkState.COMPLETED - assert agent.transfer(handle) == "ERR" - - assert endpoint.submit_attempts == 1 - assert handle.work.closed is True - assert "KeyboardInterrupt: submit return interrupted" in str( - agent.pop_xfer_error(handle) - ) - agent.remove_remote_agent("decode-agent") - agent.deregister_memory(registrations) - - -def test_recovery_precedes_hostile_exception_diagnostics_and_batch_owns_handle(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - endpoint.postcommit_error = _UnprintableInterrupt() - - def make_recovered_work_uncancellable(): - endpoint.submissions[-1][3].cancel_completes = False - - endpoint.postcommit_hook = make_recovered_work_uncancellable - handles = [] - agent.begin_handle_batch(handles) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") - - with patch.object(agent, "_start_retained_work_sweeper"): - assert handles == [handle] - assert agent.transfer(handle) == "PROC" - assert endpoint.submit_attempts == 1 - assert handle.work.closed is False - - handle.work.done = True - handle.work.state = _WorkState.COMPLETED - assert agent.check_xfer_state(handle) == "ERR" - - agent.end_handle_batch(handles) - assert handle.work.closed is True - assert "diagnostic formatting unavailable" in str(agent.pop_xfer_error(handle)) - - -def test_async_notification_lost_return_recovers_by_token(): - agent, endpoint, _ = _make_agent(synchronous_notification_send=False) - _register_and_add_peer(agent) - distractor = _Work(object()) - endpoint._open_works.append(distractor) - endpoint.postcommit_error = KeyboardInterrupt("notification return interrupted") - - with pytest.raises(KeyboardInterrupt, match="notification return interrupted"): - agent.send_notif("decode-agent", b"abort") - - _, payload, token, work = endpoint.notifications[-1] - assert payload == b"abort" - assert token.work is work - assert token.terminal_state == "ERR" - assert token.phase == "released" - assert work.closed is True - assert distractor.closed is False - assert endpoint.open_works_reads == 1 - assert agent._works == {} - assert agent._detached_works == [] - - -def test_async_notification_tracking_interruption_rolls_back_hidden_handle(): - agent, endpoint, _ = _make_agent(synchronous_notification_send=False) - _register_and_add_peer(agent) - original_track = agent._track_work - - def store_then_interrupt(handle): - original_track(handle) - raise KeyboardInterrupt("tracking return interrupted") - - with patch.object(agent, "_track_work", side_effect=store_then_interrupt): - with pytest.raises(KeyboardInterrupt, match="tracking return interrupted"): - agent.send_notif("decode-agent", b"abort") - - assert endpoint.notifications == [] - assert agent._works == {} - assert agent._detached_works == [] - - -def test_prepared_plan_teardown_quarantines_then_rejects_stale_post(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") - remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") - handle = agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) - plan = endpoint.plans[-1] - - with pytest.raises(RuntimeError, match="referenced by a transfer handle"): - agent.remove_remote_agent("decode-agent") - assert plan.closed is False - assert "decode-agent" in agent._retiring_peers - assert endpoint.submit_attempts == 0 - - with pytest.raises(RuntimeError, match="teardown is in progress"): - agent.transfer(handle) - assert endpoint.submit_attempts == 0 - assert endpoint.open_works_reads == 0 - assert handle.terminal_state == "ERR" - agent.remove_remote_agent("decode-agent") - assert plan.closed is True - assert "decode-agent" not in agent._retiring_peers - - -@pytest.mark.parametrize("version", [None, True, 1]) -def test_core_factory_v2_gate_rejects_absent_bool_and_old_versions(version): - api = _make_api() - api.BACKEND_FACTORY_API_VERSION = version - - with patch( - "sglang.srt.disaggregation.nixl.torch_transfer.importlib.import_module", - return_value=api, - ): - with pytest.raises(ImportError, match="FactoryV2"): - _load_transfer_api() - - -def test_core_factory_v2_gate_requires_marker_and_accepts_newer_core(): - api = _make_api() - del api.BackendFactoryV2 - with patch( - "sglang.srt.disaggregation.nixl.torch_transfer.importlib.import_module", - return_value=api, - ): - with pytest.raises(ImportError, match="FactoryV2"): - _load_transfer_api() - - api = _make_api() - api.BACKEND_FACTORY_API_VERSION = 3 - with patch( - "sglang.srt.disaggregation.nixl.torch_transfer.importlib.import_module", - return_value=api, - ): - assert _load_transfer_api() is api - - -@pytest.mark.parametrize("version", [None, True, 1, 3]) -def test_provider_factory_v2_gate_requires_exact_integer_abi(version): - register = MagicMock() - provider = SimpleNamespace( - TORCH_TRANSFER_FACTORY_API_VERSION=version, - register_torch_backend=register, - ) - with patch( - "sglang.srt.disaggregation.nixl.torch_transfer.importlib.import_module", - return_value=provider, - ): - with pytest.raises(ImportError, match="exact FactoryV2"): - _register_nixl_backend() - register.assert_not_called() - - provider.TORCH_TRANSFER_FACTORY_API_VERSION = 2 - with patch( - "sglang.srt.disaggregation.nixl.torch_transfer.importlib.import_module", - return_value=provider, - ): - _register_nixl_backend() - register.assert_called_once_with() - - -@pytest.mark.parametrize("mismatch", [False, True], ids=["exact", "foreign"]) -def test_constructor_recovery_is_identity_gated_and_cleanup_is_retryable(mismatch): - class HostileCleanup(RuntimeError): - def __str__(self): - raise AssertionError("cleanup diagnostics evaluated __str__") - - agent = _make_uncreated_agent() - error = KeyboardInterrupt("constructor cut") - _Endpoint.constructor_cut = ( - error, - mismatch, - ( - [AssertionError("mismatched endpoint must not close")] - if mismatch - else [HostileCleanup()] - ), - ) - - with pytest.raises(KeyboardInterrupt, match="constructor cut"): - agent.create_backend("UCX", {}) - - candidate = _Endpoint.instances[-1] - assert agent._endpoint is None - if mismatch: - assert candidate.close_calls == 0 - assert candidate.closed is False - assert agent._endpoint_recovery is None - assert any("non-matching" in note for note in error.__notes__) - candidate.close_errors.clear() - candidate.close() - else: - assert candidate.close_calls == 1 - assert candidate.closed is False - assert agent._endpoint_recovery is candidate - assert any( - "cleanup also raised:" in note and "HostileCleanup" in note - for note in error.__notes__ - ) - agent.create_backend("UCX", {}) - assert candidate.close_calls == 2 - assert candidate.closed is True - assert agent._endpoint_recovery is None - - -def test_postpublication_constructor_cut_rolls_back_all_publications(): - agent = _make_uncreated_agent() - armed = True - - def store_then_interrupt(self, name, value): - nonlocal armed - object.__setattr__(self, name, value) - if self is agent and name == "_endpoint" and type(value) is _Endpoint and armed: - armed = False - value.close_errors.append(RuntimeError("close cut")) - raise KeyboardInterrupt("after endpoint publication") - - with patch.object(TorchTransferAgent, "__setattr__", store_then_interrupt): - with pytest.raises(KeyboardInterrupt, match="after endpoint publication"): - agent.create_backend("UCX", {}) - - failed_endpoint = _Endpoint.instances[-1] - assert agent._endpoint is None - assert agent._transport_backend is None - assert agent._send_notification_sync is None - assert agent._endpoint_recovery is failed_endpoint - agent.create_backend("UCX", {}) - assert failed_endpoint.closed is True - assert agent._endpoint is _Endpoint.instances[-1] - assert agent._endpoint is not failed_endpoint - - -def test_capability_failure_rolls_back_endpoint_before_publication(): - agent = _make_uncreated_agent() - _Endpoint.missing_capability = "write" - - with pytest.raises(NotImplementedError, match="write"): - agent.create_backend("UCX", {}) - - failed_endpoint = _Endpoint.instances[-1] - assert failed_endpoint.closed is True - assert agent._endpoint is None - assert agent._transport_backend is None - assert agent._send_notification_sync is None - - -def test_registration_and_peer_lost_returns_close_only_exact_orphans(): - agent, endpoint, _ = _make_agent() - known_registration = agent.register_memory([(0x1000, 0x100, 0, "")], "VRAM") - original_register = endpoint.register - orphan_registration = None - - def register_then_interrupt(*args, **kwargs): - nonlocal orphan_registration - orphan_registration = original_register(*args, **kwargs) - raise KeyboardInterrupt("registration return cut") - - with patch.object(endpoint, "register", side_effect=register_then_interrupt): - with pytest.raises(KeyboardInterrupt, match="registration return cut"): - agent.register_memory([(0x2000, 0x100, 0, "")], "VRAM") - - assert orphan_registration.closed is True - assert known_registration[0].closed is False - assert endpoint.open_registrations == known_registration - assert agent._construction_dirty is False - - original_import = endpoint.import_peer - orphan_peer = None - - def import_then_interrupt(*args, **kwargs): - nonlocal orphan_peer - orphan_peer = original_import(*args, **kwargs) - raise KeyboardInterrupt("peer return cut") - - with patch.object(endpoint, "import_peer", side_effect=import_then_interrupt): - with pytest.raises(KeyboardInterrupt, match="peer return cut"): - agent.add_remote_agent(b"peer-metadata", peer_name="decode-agent") - - assert orphan_peer.closed is True - assert endpoint.open_peers == () - assert agent._peers == {} - assert agent._construction_dirty is False - - -@pytest.mark.parametrize("indexed", [True, False], ids=["cached", "one-off"]) -def test_plan_lost_return_preserves_known_child_and_closes_exact_orphan(indexed): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - known_local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") - known_remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") - agent.make_prepped_xfer("WRITE", known_local, [0], known_remote, [0]) - known_plan = endpoint.plans[-1] - original_prepare = endpoint.prepare - orphan_plan = None - - def prepare_then_interrupt(**kwargs): - nonlocal orphan_plan - orphan_plan = original_prepare(**kwargs) - raise KeyboardInterrupt("plan return cut") - - with patch.object(endpoint, "prepare", side_effect=prepare_then_interrupt): - with pytest.raises(KeyboardInterrupt, match="plan return cut"): - if indexed: - local = agent.prep_xfer_dlist("", [(0x1040, 32, 0)], "VRAM") - remote = agent.prep_xfer_dlist( - "decode-agent", [(0x5040, 32, 0)], "VRAM" - ) - agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) - else: - local = agent.get_xfer_descs([(0x1040, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5040, 32, 0)], "VRAM") - agent.initialize_xfer("WRITE", local, remote, "decode-agent") - - assert orphan_plan.closed is True - assert known_plan.closed is False - assert endpoint.open_plans == (known_plan,) - assert agent._construction_dirty is False - - -def test_dirty_plan_cleanup_cut_is_retried_before_next_construction(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - original_prepare = endpoint.prepare - orphan_plan = None - - def prepare_then_interrupt(**kwargs): - nonlocal orphan_plan - orphan_plan = original_prepare(**kwargs) - orphan_plan.close_errors.append(RuntimeError("first close cut")) - raise KeyboardInterrupt("plan return cut") - - with patch.object(endpoint, "prepare", side_effect=prepare_then_interrupt): - with pytest.raises(KeyboardInterrupt) as caught: - agent.initialize_xfer("WRITE", local, remote, "decode-agent") - - assert orphan_plan.closed is False - assert orphan_plan.close_calls == 1 - assert agent._construction_dirty is True - assert any("cleanup also raised" in note for note in caught.value.__notes__) - - handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") - assert orphan_plan.closed is True - assert orphan_plan.close_calls == 2 - assert handle.owned_plan is endpoint.plans[-1] - assert agent._construction_dirty is False - - -def test_reconcile_scans_future_slot_and_binding_children_first(): - agent, endpoint, _ = _make_agent() - close_order = [] - - class OrderedResource(_Resource): - def __init__(self, name): - super().__init__() - self.name = name - - def close(self): - close_order.append(self.name) - super().close() - - slot = OrderedResource("slot") - binding = OrderedResource("binding") - plan = OrderedResource("plan") - peer = OrderedResource("peer") - registration = OrderedResource("registration") - endpoint.transfer_slots.append(slot) - endpoint.bound_transfers.append(binding) - endpoint.plans.append(plan) - endpoint.peers.append(peer) - endpoint.registrations.append(registration) - agent._construction_dirty = True - - agent.register_memory([(0x1000, 0x100, 0, "")], "VRAM") - - assert slot.closed is True - assert binding.closed is True - assert plan.closed is True - assert peer.closed is True - assert registration.closed is True - assert close_order == ["slot", "binding", "plan", "peer", "registration"] - assert agent._construction_dirty is False - - -def test_clean_one_off_and_cached_hit_never_read_construction_snapshots(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - endpoint.construction_snapshot_reads = 0 - endpoint.poison_construction_snapshots = True - local_raw = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote_raw = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - agent.initialize_xfer("WRITE", local_raw, remote_raw, "decode-agent") - local = agent.prep_xfer_dlist("", [(0x1040, 32, 0)], "VRAM") - remote = agent.prep_xfer_dlist("decode-agent", [(0x5040, 32, 0)], "VRAM") - agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) - agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) - - assert endpoint.construction_snapshot_reads == 0 - - -def test_peer_postpublication_cut_retries_without_duplicate_core_child(): - class StoreThenInterrupt(dict): - def __setitem__(self, key, value): - super().__setitem__(key, value) - raise KeyboardInterrupt("peer map STORE cut") - - agent, endpoint, _ = _make_agent() - agent._peers = StoreThenInterrupt() - - with pytest.raises(KeyboardInterrupt, match="peer map STORE cut"): - agent.add_remote_agent(b"peer-metadata", peer_name="decode-agent") - - peer = endpoint.open_peers[0] - assert agent._peers["decode-agent"] is peer - assert agent._peer_wire_metadata == {"decode-agent": b"peer-metadata"} - assert agent._source_aliases["remote-id"] == "decode-agent" - assert agent._source_aliases[("remote-id", "remote-incarnation")] == "decode-agent" - assert agent._construction_dirty is False - assert ( - agent.add_remote_agent(b"peer-metadata", peer_name="decode-agent") - == "decode-agent" - ) - assert endpoint.open_peers == (peer,) - - -def test_cached_plan_postpublication_cut_retries_without_duplicate_core_child(): - class StoreThenInterrupt(dict): - def __setitem__(self, key, value): - super().__setitem__(key, value) - raise KeyboardInterrupt("plan cache STORE cut") - - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - local = agent.prep_xfer_dlist("", [(0x1000, 32, 0)], "VRAM") - remote = agent.prep_xfer_dlist("decode-agent", [(0x5000, 32, 0)], "VRAM") - agent._plan_cache = StoreThenInterrupt() - - with pytest.raises(KeyboardInterrupt, match="plan cache STORE cut"): - agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) - - plan = endpoint.open_plans[0] - cached = next(iter(agent._plan_cache.values())) - assert cached.plan is plan - assert agent._construction_dirty is False - handle = agent.make_prepped_xfer("WRITE", local, [0], remote, [0]) - assert handle.plan is plan - assert endpoint.open_plans == (plan,) - - -def test_tracking_store_cut_cannot_hide_handle_from_worker_batch(): - class StoreThenInterrupt(dict): - armed = True - - def __setitem__(self, key, value): - super().__setitem__(key, value) - if self.armed: - self.armed = False - raise KeyboardInterrupt("work store cut") - - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - agent._works = StoreThenInterrupt() - handles = [] - agent.begin_handle_batch(handles) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - - with pytest.raises(KeyboardInterrupt, match="work store cut"): - agent.initialize_xfer("WRITE", local, remote, "decode-agent") - - assert handles == [] - assert agent._works == {} - assert endpoint.plans[-1].closed is True - agent.end_handle_batch(handles) - - -def test_begin_batch_store_cut_rolls_back_both_publications(): - class StoreThenInterruptLocal: - def __init__(self): - object.__setattr__(self, "armed", True) - - def __setattr__(self, name, value): - object.__setattr__(self, name, value) - if name == "value" and self.armed: - object.__setattr__(self, "armed", False) - raise KeyboardInterrupt("batch TLS store cut") - - agent, _, _ = _make_agent() - handles = [] - agent._batch_handles = StoreThenInterruptLocal() - - with pytest.raises(KeyboardInterrupt, match="batch TLS store cut"): - agent.begin_handle_batch(handles) - - assert agent._active_handle_batches == {} - assert getattr(agent._batch_handles, "value", None) is None - - -def test_post_tracking_cleanup_cut_is_retained_and_reaped(): - agent, endpoint, _ = _make_agent() - _register_and_add_peer(agent) - handles = [] - agent.begin_handle_batch(handles) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - original_prepare = endpoint.prepare - original_track = agent._track_work - - def prepare_with_one_close_cut(**kwargs): - plan = original_prepare(**kwargs) - plan.close_errors.append(RuntimeError("plan close cut")) - return plan - - def track_then_interrupt(handle): - original_track(handle) - raise KeyboardInterrupt("tracking return cut") - - with ( - patch.object(endpoint, "prepare", side_effect=prepare_with_one_close_cut), - patch.object(agent, "_track_work", side_effect=track_then_interrupt), - patch.object(agent, "_start_retained_work_sweeper"), - ): - with pytest.raises(KeyboardInterrupt, match="tracking return cut"): - agent.initialize_xfer("WRITE", local, remote, "decode-agent") - - handle = handles[0] - assert agent._works[id(handle)] is handle - assert id(handle) in agent._retained_work_ids - assert handle.owned_plan.closed is False - agent.end_handle_batch(handles) - with agent._lock: - agent._reap_retained_works() - - assert handle.phase == "released" - assert handle.owned_plan.closed is True - assert agent._retained_work_ids == set() - - -def test_whole_batch_retention_reaps_without_per_handle_journal(): - agent, _, _ = _make_agent() - _register_and_add_peer(agent) - handles = [] - agent.begin_handle_batch(handles) - local = agent.get_xfer_descs([(0x1000, 32, 0)], "VRAM") - remote = agent.get_xfer_descs([(0x5000, 32, 0)], "VRAM") - handle = agent.initialize_xfer("WRITE", local, remote, "decode-agent") - assert handles == [handle] - assert agent.transfer(handle) == "PROC" - handle.work.cancel_completes = False - - with patch.object(agent, "_start_retained_work_sweeper"): - agent.retain_handle_batch(handles) - agent.end_handle_batch(handles) - assert agent._retained_handle_batches[id(handles)] is handles - assert id(handle) not in agent._retained_work_ids - handle.work.done = True - handle.work.state = _WorkState.COMPLETED - with agent._lock: - agent._reap_retained_works() - - assert handle.phase == "released" - assert agent._retained_handle_batches == {} - - -def test_retiring_peer_cleanup_retries_before_same_metadata_reimport(): - agent, endpoint, _ = _make_agent() - agent.add_remote_agent(b"peer-metadata", peer_name="decode-agent") - old_peer = endpoint.peers[-1] - old_peer.close_errors.append(KeyboardInterrupt("peer close cut")) - - with pytest.raises(KeyboardInterrupt, match="peer close cut"): - agent.rollback_remote_agent("decode-agent", b"peer-metadata") - - assert "decode-agent" in agent._retiring_peers - assert ( - agent.add_remote_agent(b"peer-metadata", peer_name="decode-agent") - == "decode-agent" - ) - assert old_peer.closed is True - assert endpoint.peers[-1] is not old_peer - assert "decode-agent" not in agent._retiring_peers diff --git a/test/registered/unit/managers/test_io_struct.py b/test/registered/unit/managers/test_io_struct.py index a756381a2d9a..c35ba141acfd 100644 --- a/test/registered/unit/managers/test_io_struct.py +++ b/test/registered/unit/managers/test_io_struct.py @@ -43,34 +43,6 @@ class TestTokenizedReqInputMsgpack(unittest.TestCase): - def test_transfer_generation_is_append_only_and_round_trips(self): - generation = "11" * 16 - self.assertEqual( - TokenizedGenerateReqInput.__struct_fields__[-1], "transfer_generation" - ) - request = TokenizedGenerateReqInput( - input_text="", - input_ids=array("q", [1, 2]), - input_embeds=None, - mm_inputs=None, - token_type_ids=None, - sampling_params=SamplingParams(), - return_logprob=False, - logprob_start_len=0, - top_logprobs_num=0, - token_ids_logprob=None, - stream=False, - routing_key="existing-wire-field", - cache_salt="existing-tail-field", - transfer_generation=generation, - ) - - decoded = msgpack_decode(msgpack_encode(request)) - - self.assertEqual(decoded.routing_key, "existing-wire-field") - self.assertEqual(decoded.cache_salt, "existing-tail-field") - self.assertEqual(decoded.transfer_generation, generation) - def test_rust_tokenized_generate_schema_stays_in_lockstep(self): """Compare the Rust wire declaration with the imported Python schema.""" rust_path = (