[PD] Tiny cleanup after KVReceiver refactor (#21760)
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
@@ -326,7 +326,6 @@ class CommonKVManager(BaseKVManager):
|
||||
host = self.bootstrap_host
|
||||
|
||||
bootstrap_na = NetworkAddress(host, self.bootstrap_port)
|
||||
bootstrap_server_url = bootstrap_na.to_host_port_str()
|
||||
url = f"{bootstrap_na.to_url()}/route"
|
||||
payload = {
|
||||
"attn_tp_size": self.attn_tp_size,
|
||||
@@ -477,6 +476,14 @@ class CommonKVSender(BaseKVSender):
|
||||
def failure_exception(self):
|
||||
raise Exception("Fake KVReceiver Exception")
|
||||
|
||||
def abort(self):
|
||||
self.kv_mgr.record_failure(
|
||||
self.bootstrap_room,
|
||||
"Aborted by AbortReq.",
|
||||
)
|
||||
# Explicitly set the status to failure since this request has been aborted
|
||||
self.conclude_state = KVPoll.Failed
|
||||
|
||||
|
||||
class CommonKVReceiver(BaseKVReceiver):
|
||||
_ctx = zmq.Context()
|
||||
@@ -494,15 +501,6 @@ class CommonKVReceiver(BaseKVReceiver):
|
||||
self.bootstrap_addr = bootstrap_addr
|
||||
self.kv_mgr = mgr
|
||||
self.conclude_state: Optional[KVPoll] = None
|
||||
self.bootstrap_infos = None
|
||||
self.prefill_info = None
|
||||
self.prefill_dp_rank = None
|
||||
self.target_tp_rank = None
|
||||
self.target_tp_ranks = None
|
||||
self.target_cp_ranks = None
|
||||
self.target_pp_ranks = None
|
||||
self.required_dst_info_num = None
|
||||
self.required_prefill_response_num = None
|
||||
self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room)
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping)
|
||||
|
||||
@@ -666,6 +664,14 @@ class CommonKVReceiver(BaseKVReceiver):
|
||||
def failure_exception(self):
|
||||
raise Exception("Fake KVReceiver Exception")
|
||||
|
||||
def abort(self):
|
||||
self.kv_mgr.record_failure(
|
||||
self.bootstrap_room,
|
||||
"Aborted by AbortReq.",
|
||||
)
|
||||
# Explicitly set the status to failure since this request has been aborted
|
||||
self.conclude_state = KVPoll.Failed
|
||||
|
||||
|
||||
class CommonKVBootstrapServer(BaseKVBootstrapServer):
|
||||
def __init__(self, host: str, port: int):
|
||||
|
||||
@@ -278,7 +278,7 @@ class DecodePreallocQueue:
|
||||
self.retracted_queue: List[Req] = []
|
||||
self.pending_reqs: List[DecodeRequest] = []
|
||||
self._ensure_retry_count: Dict[str, int] = {}
|
||||
self._max_ensure_retries: int = 20 # scheduling cycles
|
||||
self._max_ensure_retries: int = 15 # scheduling cycles
|
||||
self._ensure_last_attempt_time: Dict[str, float] = {}
|
||||
self._ensure_retry_interval: float = 1.0 # seconds
|
||||
self.kv_manager = self._init_kv_manager()
|
||||
@@ -546,16 +546,7 @@ class DecodePreallocQueue:
|
||||
error_msg = f"Could not fetch prefill parallel info from {bootstrap_addr} after {count} attempts"
|
||||
logger.error(error_msg)
|
||||
for decode_req in reqs:
|
||||
prepare_abort(
|
||||
decode_req.req,
|
||||
error_msg,
|
||||
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
|
||||
)
|
||||
if self.scheduler.enable_metrics:
|
||||
self.scheduler.metrics_collector.increment_bootstrap_failed_reqs()
|
||||
self.scheduler.stream_output(
|
||||
[decode_req.req], decode_req.req.return_logprob
|
||||
)
|
||||
decode_req.kv_receiver.abort()
|
||||
del self._ensure_retry_count[bootstrap_addr]
|
||||
del self._ensure_last_attempt_time[bootstrap_addr]
|
||||
else:
|
||||
|
||||
@@ -972,10 +972,8 @@ class MoriKVSender(CommonKVSender):
|
||||
raise RuntimeError(failure_reason)
|
||||
|
||||
def abort(self):
|
||||
reason = "Aborted by AbortReq."
|
||||
self.kv_mgr.record_failure(self.bootstrap_room, reason)
|
||||
self._notify_decode(KVPoll.Failed, reason)
|
||||
self.conclude_state = KVPoll.Failed
|
||||
super().abort()
|
||||
self._notify_decode(KVPoll.Failed, "Aborted by AbortReq.")
|
||||
|
||||
|
||||
class MoriKVReceiver(CommonKVReceiver):
|
||||
@@ -1106,10 +1104,8 @@ class MoriKVReceiver(CommonKVReceiver):
|
||||
def abort(self):
|
||||
if self.bootstrap_room is None:
|
||||
return
|
||||
reason = "Aborted by AbortReq."
|
||||
self.kv_mgr.record_failure(self.bootstrap_room, reason)
|
||||
super().abort()
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
|
||||
self.conclude_state = KVPoll.Failed
|
||||
self.clear()
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user