From 9a40e42384c17c130aefe85f7e6e9f357909e181 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sat, 6 Jun 2026 22:45:36 +0000 Subject: [PATCH] fix(disagg/mooncake): notify prefill on decode-side abort Port the decode->prefill half of upstream sgl-project/sglang #27372 (the worker skip-Failed guard half landed in cb1a03f0a). In multi-prefill/multi-decode PD, a decode-initiated abort frees the decode's KV pages back to the allocator, but the prefill never learns of it (request_status is per-process), so the prefill keeps RDMA-writing the remaining chunks into pages the decode may have reallocated to a different live request -> KV corruption. The already-ported worker guard is a no-op here because nothing sets the prefill room to Failed. Now the decode receiver sends a 4-field b"ABORT" notification to every prefill peer (on abort() and on the poll() WaitingForInput timeout, at most once); the prefill marks the room Failed (so the worker guard fires) and replies b"ABORT_ACK". The ABORT branch is handled before the unconditional waiting_req_bytes[3] decode in bootstrap_thread (a 4-field ABORT would otherwise crash the else branch), and ABORT_ACK before the 3-tuple unpack in the decode thread. De-entangled from the staging/tracing infra this branch lacks. Component-tested (test_mooncake_abort_protocol.py: format, send-once guard, abort wiring, ZMQ round-trip). Cross-process corruption-window closure to be verified by the multi-P/D PD harness. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../srt/disaggregation/mooncake/conn.py | 68 +++++++++ .../test_mooncake_abort_protocol.py | 138 ++++++++++++++++++ 2 files changed, 206 insertions(+) create mode 100644 test/registered/unit/disaggregation/test_mooncake_abort_protocol.py diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 0e08c1cb3..2c0aa805a 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -1224,6 +1224,34 @@ class MooncakeKVManager(CommonKVManager): while True: waiting_req_bytes = self.server_socket.recv_multipart() room = waiting_req_bytes[0].decode("ascii") + # Decode-side abort notification: mark the room Failed and ACK, so the + # transfer worker (which skips Failed rooms) stops RDMA-writing into KV + # pages the decode has freed/reused. Port of upstream sgl-project/sglang + # #27372 (decode->prefill half). MUST be handled here, before the + # UNCONDITIONAL waiting_req_bytes[3] decode below: a 4-field ABORT + # message [ABORT, room, ip, port] would otherwise fall through into the + # else branch and crash on waiting_req_bytes[7]. + 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")) + if ( + room_to_be_aborted in self.request_status + and self.check_status(room_to_be_aborted) != KVPoll.Success + ): + self.update_status(room_to_be_aborted, KVPoll.Failed) + try: + na = NetworkAddress(decode_ip, decode_port) + self._connect( + na.to_tcp(), is_ipv6=na.is_ipv6 + ).send_multipart( + [b"ABORT_ACK", str(room_to_be_aborted).encode("ascii")] + ) + 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") if room == "None": self.decode_kv_args_table[mooncake_session_id] = ( @@ -1261,6 +1289,12 @@ class MooncakeKVManager(CommonKVManager): self._handle_aux_data(msg) continue + if msg[0] == b"ABORT_ACK": + # Prefill acknowledged our abort; the room is already Failed locally, + # so nothing more to do. Port of upstream #27372. Must precede the + # 3-tuple unpack below (a 2-field ACK would otherwise raise). + continue + bootstrap_room, status, prefill_rank = msg status = int(status.decode("ascii")) bootstrap_room = int(bootstrap_room.decode("ascii")) @@ -1617,6 +1651,7 @@ class MooncakeKVReceiver(CommonKVReceiver): self.session_id = mgr.get_session_id() self.conclude_state = None self.init_time = None + self.abort_notified = False super().__init__(mgr, bootstrap_addr, bootstrap_room, prefill_dp_rank) self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room) @@ -1746,6 +1781,8 @@ class MooncakeKVReceiver(CommonKVReceiver): self.bootstrap_room, f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s in KVPoll.WaitingForInput", ) + # Tell prefill peers to stop transferring into our pages. + self._send_abort_notification() self.conclude_state = KVPoll.Failed return KVPoll.Failed @@ -1777,6 +1814,35 @@ class MooncakeKVReceiver(CommonKVReceiver): ) raise KVTransferError(self.bootstrap_room, failure_reason) + def _send_abort_notification(self): + # Notify every prefill peer that this room is aborting, so the prefill stops + # RDMA-writing into KV pages this decode is about to free/reuse. Without this, + # the prefill never learns of a decode-initiated abort (its request_status is + # per-process) and the worker's skip-Failed guard never fires. Port of upstream + # sgl-project/sglang #27372 (decode->prefill half). Sent at most once. + if self.abort_notified: + return + bootstrap_infos = getattr(self, "bootstrap_infos", None) + if not bootstrap_infos: + return + self.abort_notified = True + for bootstrap_info in bootstrap_infos: + try: + sock, lock = self._connect_to_bootstrap_server(bootstrap_info) + with lock: + sock.send_multipart( + [ + b"ABORT", + str(self.bootstrap_room).encode("ascii"), + self.kv_mgr.local_ip.encode("ascii"), + str(self.kv_mgr.rank_port).encode("ascii"), + ] + ) + except Exception as e: + logger.debug( + f"Failed to send ABORT for room {self.bootstrap_room}: {e}" + ) + def abort(self): self.kv_mgr.record_failure( self.bootstrap_room, @@ -1786,6 +1852,8 @@ class MooncakeKVReceiver(CommonKVReceiver): # transfer worker and other pollers observe the abort. Port of upstream # sgl-project/sglang #24522. self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) + # Tell the prefill peers to stop transferring into our (soon-freed) pages. + self._send_abort_notification() # Explicitly set the status to failure since this request has been aborted self.conclude_state = KVPoll.Failed diff --git a/test/registered/unit/disaggregation/test_mooncake_abort_protocol.py b/test/registered/unit/disaggregation/test_mooncake_abort_protocol.py new file mode 100644 index 000000000..3cdf42a28 --- /dev/null +++ b/test/registered/unit/disaggregation/test_mooncake_abort_protocol.py @@ -0,0 +1,138 @@ +"""Component tests for the decode->prefill abort protocol (port of upstream #27372). + +These cover the pieces that do NOT need a live PD transfer: + - the decode side builds the exact 4-field b"ABORT" message and sends it at most once, + - MooncakeKVReceiver.abort() invokes the notification and marks the room Failed, + - the ZMQ wire format round-trips so the prefill handler's waiting_req_bytes[1..3] + indices are always present (the branch-ordering trap: a 4-field ABORT must not + IndexError before the room=="ABORT" check). + +The actual cross-process corruption-window closure is verified by the multi-P/D PD +harness, not here. +""" + +import unittest + +from sglang.srt.disaggregation.base.conn import KVPoll +from sglang.srt.disaggregation.mooncake.conn import MooncakeKVReceiver + + +class _DummyLock: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + +class _CapturingSock: + def __init__(self): + self.sent = [] + + def send_multipart(self, parts): + self.sent.append(parts) + + +class _FakeMgr: + def __init__(self): + self.local_ip = "10.0.0.1" + self.rank_port = 12345 + self.status_updates = [] + + def record_failure(self, room, reason): + pass + + def update_status(self, room, status): + self.status_updates.append((room, status)) + + +def _make_receiver(bootstrap_infos): + r = MooncakeKVReceiver.__new__(MooncakeKVReceiver) # bypass heavy __init__ + r.abort_notified = False + r.bootstrap_room = 777 + r.bootstrap_infos = bootstrap_infos + r.kv_mgr = _FakeMgr() + r.conclude_state = None + return r + + +class TestSendAbortNotification(unittest.TestCase): + def test_message_format_is_four_fields(self): + sock = _CapturingSock() + r = _make_receiver([{"peer": 1}]) + r._connect_to_bootstrap_server = lambda info: (sock, _DummyLock()) + r._send_abort_notification() + self.assertEqual(len(sock.sent), 1) + parts = sock.sent[0] + # [ABORT, room, ip, port] — exactly 4 fields so the prefill handler's + # waiting_req_bytes[3] (port) access never IndexErrors. + self.assertEqual(len(parts), 4) + self.assertEqual(parts[0], b"ABORT") + self.assertEqual(parts[1], b"777") + self.assertEqual(parts[2], b"10.0.0.1") + self.assertEqual(parts[3], b"12345") + + def test_sent_at_most_once(self): + sock = _CapturingSock() + r = _make_receiver([{"peer": 1}]) + r._connect_to_bootstrap_server = lambda info: (sock, _DummyLock()) + r._send_abort_notification() + r._send_abort_notification() + self.assertEqual(len(sock.sent), 1) + self.assertTrue(r.abort_notified) + + def test_no_bootstrap_infos_is_noop(self): + r = _make_receiver(None) + r._send_abort_notification() # must not raise + self.assertFalse(r.abort_notified) + + def test_sends_to_every_peer(self): + socks = [_CapturingSock(), _CapturingSock()] + infos = [{"peer": 0}, {"peer": 1}] + r = _make_receiver(infos) + r._connect_to_bootstrap_server = lambda info: (socks[info["peer"]], _DummyLock()) + r._send_abort_notification() + self.assertEqual(len(socks[0].sent), 1) + self.assertEqual(len(socks[1].sent), 1) + + +class TestAbortWiring(unittest.TestCase): + def test_abort_notifies_and_marks_failed(self): + r = _make_receiver([{"peer": 1}]) + called = [] + r._send_abort_notification = lambda: called.append(True) + r.abort() + self.assertTrue(called) + self.assertEqual(r.conclude_state, KVPoll.Failed) + self.assertIn((777, KVPoll.Failed), r.kv_mgr.status_updates) + + +class TestZmqWireFormat(unittest.TestCase): + def test_abort_and_ack_roundtrip(self): + import zmq + + ctx = zmq.Context.instance() + pull = ctx.socket(zmq.PULL) + pull.bind("tcp://127.0.0.1:*") + ep = pull.getsockopt_string(zmq.LAST_ENDPOINT) + push = ctx.socket(zmq.PUSH) + push.connect(ep) + try: + push.send_multipart([b"ABORT", b"777", b"10.0.0.1", b"12345"]) + parts = pull.recv_multipart() + self.assertEqual(parts[0], b"ABORT") + self.assertEqual(int(parts[1].decode()), 777) + self.assertEqual(parts[2].decode(), "10.0.0.1") + self.assertEqual(int(parts[3].decode()), 12345) + + push.send_multipart([b"ABORT_ACK", b"777"]) + ack = pull.recv_multipart() + self.assertEqual(ack[0], b"ABORT_ACK") + self.assertEqual(int(ack[1].decode()), 777) + finally: + push.close() + pull.close() + + +if __name__ == "__main__": + unittest.main()