From f04c522534e65417bbbce2a808740d3ec4f7540e Mon Sep 17 00:00:00 2001
From: Shangming Cai <csmthu@gmail.com>
Date: Mon, 18 May 2026 19:56:42 +0800
Subject: [PATCH] [PD] Add conclude_state to fake KV backend (#25599)

Signed-off-by: Shangming Cai <csmthu@gmail.com>
---
 python/sglang/srt/disaggregation/fake/conn.py | 24 +++++++++++++++----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/python/sglang/srt/disaggregation/fake/conn.py b/python/sglang/srt/disaggregation/fake/conn.py
index e44cbb7b3..09c5f80ce 100644
--- a/python/sglang/srt/disaggregation/fake/conn.py
+++ b/python/sglang/srt/disaggregation/fake/conn.py
@@ -46,15 +46,19 @@ class FakeKVSender(BaseKVSender):
     ):
         self.kv_mgr = mgr
         self.has_sent = False
+        self.conclude_state: Optional[KVPoll] = None
 
     def poll(self) -> KVPoll:
-        if self.has_sent is False:
+        if self.conclude_state is not None:
+            return self.conclude_state
+        if not self.has_sent:
             # Assume handshake completed instantly
             return KVPoll.WaitingForInput
-        else:
-            # Assume transfer completed instantly
-            logger.debug("FakeKVSender poll success")
-            return KVPoll.Success
+
+        # Assume transfer completed instantly
+        logger.debug("FakeKVSender poll success")
+        self.conclude_state = KVPoll.Success
+        return KVPoll.Success
 
     def get_transfer_metric(self) -> KVTransferMetric:
         return KVTransferMetric()
@@ -82,6 +86,9 @@ class FakeKVSender(BaseKVSender):
     def failure_exception(self):
         raise Exception("Fake KVSender Exception")
 
+    def abort(self):
+        self.conclude_state = KVPoll.Failed
+
 
 class FakeKVReceiver(BaseKVReceiver):
     def __init__(
@@ -93,13 +100,17 @@ class FakeKVReceiver(BaseKVReceiver):
         self.bootstrap_done = False
         self.has_sent_metadata = False
         self.require_staging: bool = False
+        self.conclude_state: Optional[KVPoll] = None
 
     def poll(self) -> KVPoll:
+        if self.conclude_state is not None:
+            return self.conclude_state
         if not self.bootstrap_done:
             return KVPoll.Bootstrapping
         if not self.has_sent_metadata:
             return KVPoll.WaitingForInput
         logger.debug("FakeKVReceiver poll success")
+        self.conclude_state = KVPoll.Success
         return KVPoll.Success
 
     def init(
@@ -122,3 +133,6 @@ class FakeKVReceiver(BaseKVReceiver):
 
     def failure_exception(self):
         raise Exception("Fake KVReceiver Exception")
+
+    def abort(self):
+        self.conclude_state = KVPoll.Failed