[PD] Add conclude_state to fake KV backend (#25599)
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user