Skip to content

[PD] Preserve abort ACKs until in-flight KV transfers drain #40645

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ShangmingCai merged 1 commit into from
Sep 22, 2026
Merged

[PD] Preserve abort ACKs until in-flight KV transfers drain #40645

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions python/sglang/srt/disaggregation/common/conn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
8 changes: 5 additions & 3 deletions python/sglang/srt/disaggregation/mooncake/conn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
):

@ShangmingCai ShangmingCai Sep 22, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: with this new condition, the elif at L2438 is only reached when the room is inactive and outstanding <= 0, so its == 0 re-check can never usefully be false. The one state where it differs is a hypothetically negative counter, where the elif silently drops the abort (no ack, no registration → decode waits out the release timeout), while a plain else would correctly ack — a negative counter still means nothing is in flight.

-                        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.

Purely cosmetic/robustness — fine to take as a follow-up; the current code is behaviorally identical in every reachable state.

self.update_status(room_to_be_aborted, KVPoll.Failed)
self.register_deferred_ack_target(
room_to_be_aborted, decode_ip, decode_port
Expand Down
20 changes: 20 additions & 0 deletions test/registered/unit/disaggregation/test_disaggregation_wire.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
from sglang.srt.disaggregation.mooncake.conn import (
KVArgsRegisterInfo,
MooncakeKVManager,
MooncakeKVSender,
TransferInfo,
)
from sglang.srt.disaggregation.utils import (
Expand Down Expand Up @@ -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",
Expand Down
Loading