diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index a140543c7..3dd6f05fd 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -1596,9 +1596,30 @@ class MooncakeKVManager(CommonKVManager): return if bootstrap_room not in self.transfer_infos: - # This means that the current rank is a dummy rank for this request, - # and it has already been marked as success, so there is no need to - # add further chunks into the transfer queue. + if self.check_status(bootstrap_room) == KVPoll.Success: + # Dummy rank for this request: it was already marked Success + # at handshake time, so there is nothing left to transfer. + return + # Non-Success room with no transfer destinations is an anomaly + # (decode peer torn down between handshake and this send, abort + # notification lost). Silently dropping the chunk used to wedge + # the request in WaitingForInput FOREVER — Success can never be + # set without the last chunk, and nothing else fails the room. + # Conclude it loudly instead; the scheduler's poll consensus + # (MIN-reduce, Failed=0) reaps it on every rank. + logger.warning( + "KV chunk for bootstrap_room=%s has no transfer destinations " + "(status=%s); failing the request instead of dropping the " + "chunk silently.", + bootstrap_room, + self.check_status(bootstrap_room), + ) + self.record_failure( + bootstrap_room, + f"Request {bootstrap_room} lost its transfer destinations " + "before the KV chunk was submitted (decode peer gone?)", + ) + self.update_status(bootstrap_room, KVPoll.Failed) return # NOTE(shangming): sharding according to the dst_infos to make sure diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 7e7b63cec..9892392e2 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -972,6 +972,11 @@ class SchedulerDisaggregationPrefillMixin: req.output_ids.append(next_token_id) self.tree_cache.cache_unfinished_req(req) # update the tree and lock self.disagg_prefill_inflight_queue.append(req) + # Residency clock for the inflight liveness bound (every rank + # stamps at the same logical point in the same tick, so the + # timeout decision below stays effectively rank-uniform; the + # MIN-reduced poll consensus absorbs any microsecond skew). + req.disagg_inflight_enter_time = time.perf_counter() if spec_cpu_row is not None: k = spec_cpu_row[i] req.output_topk_p = spec_topk_p_cpu[k] @@ -1102,6 +1107,34 @@ class SchedulerDisaggregationPrefillMixin: continue if poll in [KVPoll.WaitingForInput, KVPoll.Transferring]: + # Liveness bound: an inflight entry whose transfer never + # concludes (decode peer torn down with its abort + # notification lost, last chunk silently dropped, decode + # crash) would otherwise be re-queued FOREVER — it pins the + # request's KV pages and rides every poll collective. On + # expiry, conclude the sender locally and keep the entry + # undone: Failed(0) wins the CP/TP MIN-reduce, so every rank + # reaps it together through the Failed branch on the NEXT + # poll — one removal path, rank-uniform by construction. + enter_time = getattr(req, "disagg_inflight_enter_time", None) + if ( + enter_time is not None + and time.perf_counter() - enter_time + >= envs.SGLANG_DISAGGREGATION_INFLIGHT_TIMEOUT.get() + and hasattr(req.disagg_kv_sender, "abort") + ): + logger.warning( + "Prefill inflight transfer timed out; failing the " + "request. rid=%s bootstrap_room=%s poll=%s " + "elapsed=%.1fs (raise " + "SGLANG_DISAGGREGATION_INFLIGHT_TIMEOUT if transfers " + "can legitimately take longer)", + req.rid, + req.bootstrap_room, + poll, + time.perf_counter() - enter_time, + ) + req.disagg_kv_sender.abort() undone_reqs.append(req) elif poll == KVPoll.Success: # transfer done release_kv_cache(req, self.tree_cache) # unlock the tree diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 1b33ff455..a27a2ace9 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -296,6 +296,10 @@ class Envs: SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL = EnvFloat(5.0) SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE = EnvInt(2) SGLANG_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300) + # Liveness bound for entries in the prefill inflight queue (transfer + # submitted, waiting to conclude). 300 matches the bootstrap/waiting + # timeouts; transfers legitimately take seconds, so do not go below ~60. + SGLANG_DISAGGREGATION_INFLIGHT_TIMEOUT = EnvInt(300) SGLANG_DISAGGREGATION_NIXL_BACKEND = EnvStr("UCX") SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER = EnvBool(False) SGLANG_DISAGGREGATION_TRANSFER_STATS = EnvBool(False) diff --git a/test/registered/unit/disaggregation/test_prefill_inflight_liveness.py b/test/registered/unit/disaggregation/test_prefill_inflight_liveness.py new file mode 100644 index 000000000..a3690a4ca --- /dev/null +++ b/test/registered/unit/disaggregation/test_prefill_inflight_liveness.py @@ -0,0 +1,85 @@ +"""Unit tests for the PD prefill inflight liveness fixes. + +Background (2026-06-12 e2e incident): a request whose decode peer was torn +down between the handshake and the prefill's final ``send()`` wedged FOREVER +in ``disagg_prefill_inflight_queue`` — ``add_transfer_request`` silently +dropped the last chunk (no transfer destinations), so Success was never set, +and the lost decode ABORT notification meant Failed was never injected. + +Registered: CPU CI. +""" + +import threading +import unittest + +import numpy as np + +from sglang.srt.disaggregation.base.conn import KVPoll +from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="stage-a-test-cpu") + +ROOM = 4242 + + +def _make_manager(status: KVPoll) -> MooncakeKVManager: + m = MooncakeKVManager.__new__(MooncakeKVManager) # skip heavy init + m.disaggregation_mode = DisaggregationMode.PREFILL + m.request_status = {ROOM: status} + m.transfer_infos = {} # no destinations registered for ROOM + m.failure_records = {} + m.failure_lock = threading.Lock() + m.transfer_queues = [] # must never be reached in these tests + return m + + +class TestAddTransferRequestWithoutDestinations(unittest.TestCase): + def _send_last_chunk(self, manager: MooncakeKVManager) -> None: + manager.add_transfer_request( + ROOM, + kv_indices=np.array([0, 1], dtype=np.int32), + index_slice=slice(0, 2), + is_last_chunk=True, + aux_index=0, + ) + + def test_non_success_room_is_failed_loudly(self): + # The wedge case: room bootstrapped (WaitingForInput) but the decode + # peer is gone. The chunk must NOT be dropped silently — the room + # must conclude Failed so the scheduler's poll consensus reaps it. + manager = _make_manager(KVPoll.WaitingForInput) + self._send_last_chunk(manager) + self.assertEqual(manager.check_status(ROOM), KVPoll.Failed) + self.assertIn(ROOM, manager.failure_records) + + def test_dummy_rank_success_room_is_untouched(self): + # Dummy ranks are marked Success at handshake time and legitimately + # have no transfer_infos: the early return must stay a no-op. + manager = _make_manager(KVPoll.Success) + self._send_last_chunk(manager) + self.assertEqual(manager.check_status(ROOM), KVPoll.Success) + self.assertNotIn(ROOM, manager.failure_records) + + def test_already_failed_room_stays_silent(self): + # Pre-existing behavior: chunks for failed rooms are dropped. + manager = _make_manager(KVPoll.Failed) + self._send_last_chunk(manager) + self.assertEqual(manager.check_status(ROOM), KVPoll.Failed) + self.assertNotIn(ROOM, manager.failure_records) + + +class TestInflightResidencyTimeoutEnv(unittest.TestCase): + def test_default_matches_sibling_disagg_timeouts(self): + from sglang.srt.environ import envs + + self.assertEqual(envs.SGLANG_DISAGGREGATION_INFLIGHT_TIMEOUT.get(), 300) + self.assertEqual( + envs.SGLANG_DISAGGREGATION_INFLIGHT_TIMEOUT.get(), + envs.SGLANG_DISAGGREGATION_WAITING_TIMEOUT.get(), + ) + + +if __name__ == "__main__": + unittest.main()