minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[PD] Add conclude_state to fake KV backend (#25599)

Browse Source 
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Shangming Cai
2026-05-18 19:56:42 +08:00
committed by GitHub
parent be64d875ad
commit f04c522534
1 changed files with 19 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/fake/conn.py
+19 -5
View File
@@ -46,15 +46,19 @@ class FakeKVSender(BaseKVSender):
): ):
self.kv_mgr = mgr self.kv_mgr = mgr
self.has_sent = False self.has_sent = False
self.conclude_state: Optional[KVPoll] = None
def poll(self) -> KVPoll: 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 # Assume handshake completed instantly
return KVPoll.WaitingForInput return KVPoll.WaitingForInput
else:
# Assume transfer completed instantly # Assume transfer completed instantly
logger.debug("FakeKVSender poll success") logger.debug("FakeKVSender poll success")
return KVPoll.Success self.conclude_state = KVPoll.Success
return KVPoll.Success
def get_transfer_metric(self) -> KVTransferMetric: def get_transfer_metric(self) -> KVTransferMetric:
return KVTransferMetric() return KVTransferMetric()
@@ -82,6 +86,9 @@ class FakeKVSender(BaseKVSender):
def failure_exception(self): def failure_exception(self):
raise Exception("Fake KVSender Exception") raise Exception("Fake KVSender Exception")
def abort(self):
self.conclude_state = KVPoll.Failed
class FakeKVReceiver(BaseKVReceiver): class FakeKVReceiver(BaseKVReceiver):
def __init__( def __init__(
@@ -93,13 +100,17 @@ class FakeKVReceiver(BaseKVReceiver):
self.bootstrap_done = False self.bootstrap_done = False
self.has_sent_metadata = False self.has_sent_metadata = False
self.require_staging: bool = False self.require_staging: bool = False
self.conclude_state: Optional[KVPoll] = None
def poll(self) -> KVPoll: def poll(self) -> KVPoll:
if self.conclude_state is not None:
return self.conclude_state
if not self.bootstrap_done: if not self.bootstrap_done:
return KVPoll.Bootstrapping return KVPoll.Bootstrapping
if not self.has_sent_metadata: if not self.has_sent_metadata:
return KVPoll.WaitingForInput return KVPoll.WaitingForInput
logger.debug("FakeKVReceiver poll success") logger.debug("FakeKVReceiver poll success")
self.conclude_state = KVPoll.Success
return KVPoll.Success return KVPoll.Success
def init( def init(
@@ -122,3 +133,6 @@ class FakeKVReceiver(BaseKVReceiver):
def failure_exception(self): def failure_exception(self):
raise Exception("Fake KVReceiver Exception") raise Exception("Fake KVReceiver Exception")
def abort(self):
self.conclude_state = KVPoll.Failed