[PD] Tiny cleanup after KVReceiver refactor (#21760)

Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Shangming Cai
2026-03-31 21:07:57 +08:00
committed by GitHub
parent c7adca9992
commit ca2b2130ba
3 changed files with 21 additions and 28 deletions
+16 -10
View File
@@ -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):
+2 -11
View File
@@ -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()