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

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
+19 -5
View File
@@ -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