From 83406b6f4501cdfb077de8760db7fca9a99c12fa Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Mon, 21 Sep 2026 16:19:57 -0700 Subject: [PATCH 1/2] Fix PD abort acknowledgments racing sender cleanup --- .../sglang/srt/disaggregation/common/conn.py | 9 ++++++--- .../srt/disaggregation/mooncake/conn.py | 8 +++++--- .../test_disaggregation_wire.py | 20 +++++++++++++++++++ 3 files changed, 31 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 526edf2ace69..3500fe2fa576 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -1654,9 +1654,12 @@ def clear(self) -> 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) + if hasattr(self.kv_mgr, "_staging_outstanding"): + # Preserve the target until in-flight writes drain, even when + # the scheduler has already observed Failed and cleared the room. + self.kv_mgr._maybe_ack_drained_abort(self.bootstrap_room) + else: + self.kv_mgr._deferred_ack_targets.pop(self.bootstrap_room, None) def abort(self): self.kv_mgr.record_failure( diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 25cbf4c80c65..ed3defa94c96 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -2419,9 +2419,11 @@ def bootstrap_thread(): # 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: + # acks once its in-flight write drains. + if ( + room_active + or self._staging_outstanding.get(room_to_be_aborted, 0) > 0 + ): self.update_status(room_to_be_aborted, KVPoll.Failed) self.register_deferred_ack_target( room_to_be_aborted, decode_ip, decode_port diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 0a9bc35d4eeb..e31eaab4e712 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -30,6 +30,7 @@ from sglang.srt.disaggregation.mooncake.conn import ( KVArgsRegisterInfo, MooncakeKVManager, + MooncakeKVSender, TransferInfo, ) from sglang.srt.disaggregation.utils import ( @@ -68,6 +69,25 @@ class TestDisaggregationWire(unittest.TestCase): + def test_sender_clear_keeps_abort_ack_until_writes_drain(self): + manager = object.__new__(MooncakeKVManager) + sender = object.__new__(MooncakeKVSender) + sender.kv_mgr, sender.bootstrap_room = manager, 42 + for outstanding in (0, 1): + with self.subTest(outstanding=outstanding): + manager.request_status = {42: KVPoll.Failed} + manager._staging_outstanding = {42: outstanding} + manager._deferred_ack_targets = {42: ("127.0.0.1", 1234)} + with patch.object(manager, "_send_abort_ack") as ack: + sender.clear() + if outstanding: + ack.assert_not_called() + manager._staging_outstanding[42] = 0 + manager._maybe_ack_drained_abort(42) + ack.assert_called_once_with("127.0.0.1", 1234, 42) + manager._maybe_ack_drained_abort(42) + ack.assert_called_once() + def test_mooncake_registration_staging_fields(self): msg = [ b"room", From 4ac5d7b8aacd46d26b6d911014f4888857c8fa1e Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Tue, 22 Sep 2026 15:15:15 +0800 Subject: [PATCH 2/2] [PD] Simplify late-abort quiescent ack branch to else With the guarding condition now taking every case with writes outstanding (room active or outstanding > 0), the elif's == 0 re-check can only differ on a negative counter, where it would silently drop the abort (no ack, no registration) and leave decode to wait out the release timeout. A negative count still means nothing is in flight, so ack unconditionally. Follow-up to #40645. Co-Authored-By: Claude Fable 5 --- python/sglang/srt/disaggregation/mooncake/conn.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index ed3defa94c96..1bc32766b657 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -2435,10 +2435,10 @@ def bootstrap_thread(): 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. + else: + # Concluded/unknown AND quiescent (the branch above + # already took every case with writes outstanding): + # ack now so decode releases without the timeout. self._send_abort_ack( decode_ip, decode_port, room_to_be_aborted )