From 9fc60a8afd64b135342a7313bf611285aeab7208 Mon Sep 17 00:00:00 2001
From: Niko Ma <nima@amd.com>
Date: Fri, 28 Aug 2026 14:02:35 +0800
Subject: [PATCH] [PD] Fix MORI-IO ABORT bootstrap message handling (#29133)

Co-authored-by: billishyahao <bill.he@amd.com>
Co-authored-by: TianDi101 <tiandi950321@163.com>
Co-authored-by: Duyi-Wang <duyi.wang@amd.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
---
 python/sglang/srt/disaggregation/mori/conn.py | 55 +++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py
index 6452fc7d3..208fcadb2 100644
--- a/python/sglang/srt/disaggregation/mori/conn.py
+++ b/python/sglang/srt/disaggregation/mori/conn.py
@@ -49,6 +49,7 @@ from sglang.srt.utils.network import NetworkAddress, get_local_ip_auto
 
 logger = logging.getLogger(__name__)
 MORI_GUARD = b"MoriMsgGuard"
+_TAG_ABORT = b"ABORT"
 
 
 def _normalize_state_indices_per_component(
@@ -705,11 +706,52 @@ class MoriKVManager(CommonKVManager):
             return None
         return payload
 
+    def _handle_abort_message(self, msg: List[bytes]) -> None:
+        """Handle best-effort ABORT notifications from the decode side."""
+        if len(msg) < 2:
+            logger.warning("Malformed ABORT message: too few frames (%d)", len(msg))
+            return
+
+        try:
+            bootstrap_room = int(msg[1].decode("ascii"))
+        except (ValueError, UnicodeDecodeError):
+            logger.warning("Malformed ABORT message: invalid room field %r", msg[1])
+            return
+
+        with self.transfer_lock:
+            current = self.request_status.get(bootstrap_room)
+            if current is None:
+                logger.debug(
+                    "ABORT for room %s is not tracked; ignoring",
+                    bootstrap_room,
+                )
+                return
+            if current == KVPoll.Success:
+                logger.debug(
+                    "ABORT for room %s already succeeded; ignoring",
+                    bootstrap_room,
+                )
+                return
+            if current == KVPoll.Failed:
+                return
+
+            self.update_status(bootstrap_room, KVPoll.Failed)
+
+        logger.debug("Room %s marked Failed via ABORT from decode", bootstrap_room)
+
     def _start_bootstrap_thread(self) -> None:
         def bootstrap_worker():
             while True:
                 try:
                     msg = self.server_socket.recv_multipart()
+                    if not msg:
+                        continue
+
+                    tag = msg[0]
+                    if tag == _TAG_ABORT:
+                        self._handle_abort_message(msg)
+                        continue
+
                     payload = self._validate_message(msg)
                     if payload is None:
                         continue
@@ -752,6 +794,12 @@ class MoriKVManager(CommonKVManager):
                         logger.warning("Incomplete status payload received")
                         continue
                     bootstrap_room = int(payload[0].decode("ascii"))
+                    if bootstrap_room not in self.request_status:
+                        logger.debug(
+                            "Dropping late status for cleared room %s",
+                            bootstrap_room,
+                        )
+                        continue
                     status_code = int(payload[1].decode("ascii"))
                     prefill_rank = int(payload[2].decode("ascii"))
                     failure_reason = (
@@ -1471,6 +1519,10 @@ class MoriKVManager(CommonKVManager):
         targets: List[TransferTarget] = []
         target_infos_snapshot: Optional[List[TransferInfo]] = None
         with self.transfer_lock:
+            current = self.request_status.get(bootstrap_room)
+            if current is None or current == KVPoll.Failed:
+                return [], None
+
             transfer_infos = self.transfer_infos.get(bootstrap_room)
             if not transfer_infos:
                 raise RuntimeError(
@@ -1805,8 +1857,11 @@ class MoriKVReceiver(CommonKVReceiver):
     def abort(self):
         if self.bootstrap_room is None:
             return
+        bootstrap_room = self.bootstrap_room
         super().abort()
         self.clear()
+        with self.kv_mgr.failure_lock:
+            self.kv_mgr.failure_records.pop(bootstrap_room, None)
 
 
 class MoriKVBootstrapServer(CommonKVBootstrapServer):