From 738cbde9027de2a2671167102af0dc41b10bf0fa Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Mon, 16 Mar 2026 14:12:13 +0800 Subject: [PATCH] [PD] Make pending reqs resolving more robust (#20505) Signed-off-by: Shangming Cai --- .../sglang/srt/disaggregation/common/conn.py | 21 +++-------- python/sglang/srt/disaggregation/decode.py | 37 +++++++++++++------ 2 files changed, 31 insertions(+), 27 deletions(-) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 32fd4dfa0..1e38ba747 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -183,26 +183,13 @@ class CommonKVManager(BaseKVManager): with self.failure_lock: self.failure_records[bootstrap_room] = failure_reason - def ensure_parallel_info( - self, bootstrap_addr: str, max_retries: int = 5, retry_interval: float = 1.0 - ) -> bool: + def ensure_parallel_info(self, bootstrap_addr: str) -> bool: """Fetch and cache prefill parallel info if not yet available. Returns True if info is available (cached or freshly fetched). - Retries with backoff if the prefill server hasn't registered yet. """ if bootstrap_addr in self.prefill_info_table: return True - info = None - for attempt in range(max_retries): - info = self._fetch_prefill_server_info(bootstrap_addr) - if info is not None: - break - if attempt < max_retries - 1: - logger.info( - f"Prefill server info not available from {bootstrap_addr}, " - f"retrying ({attempt + 1}/{max_retries})..." - ) - time.sleep(retry_interval) + info = self._fetch_prefill_server_info(bootstrap_addr) if info is None: return False @@ -427,10 +414,12 @@ class CommonKVReceiver(BaseKVReceiver): self.kv_mgr = mgr self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping) - if not self.kv_mgr.ensure_parallel_info(self.bootstrap_addr): + if self.bootstrap_addr not in self.kv_mgr.prefill_info_table: self.kv_mgr.record_failure( self.bootstrap_room, f"Could not fetch prefill parallel info from bootstrap_addr: {self.bootstrap_addr}", + f"Prefill server with bootstrap_addr: {self.bootstrap_addr} is down, prefill_info_table has been cleared." + f"Request (bootstrap_addr: {self.bootstrap_addr}) has been marked as failed.", ) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) self.bootstrap_infos = None diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 9d55ba027..36e99b28b 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -265,6 +265,7 @@ class DecodePreallocQueue: self.queue: List[DecodeRequest] = [] self.retracted_queue: List[Req] = [] self.pending_reqs: List[Req] = [] + self._bootstrap_addr_retry_count: Dict[str, int] = {} self.kv_manager = self._init_kv_manager() if self.scheduler.tp_worker.is_hybrid_swa: @@ -513,19 +514,33 @@ class DecodePreallocQueue: # which is a conflict with the lazy resolve logic in CommonKVReceiver, # so we need to ensure the parallel info before resolving it. if not self.kv_manager.ensure_parallel_info(bootstrap_addr): - error_message = f"Could not fetch prefill parallel info from bootstrap server {bootstrap_addr}" - logger.error(error_message) - for req in reqs: - prepare_abort( - req, - error_message, - status_code=HTTPStatus.INTERNAL_SERVER_ERROR, - ) - if self.scheduler.enable_metrics: - self.scheduler.metrics_collector.increment_bootstrap_failed_reqs() - self.scheduler.stream_output([req], req.return_logprob) + # Increment retry count for this bootstrap_addr + retry_count = ( + self._bootstrap_addr_retry_count.get(bootstrap_addr, 0) + 1 + ) + self._bootstrap_addr_retry_count[bootstrap_addr] = retry_count + + if retry_count >= 5: + # Max retries reached, fail the requests + error_message = f"Could not fetch prefill parallel info from bootstrap server {bootstrap_addr} after {retry_count} retries" + logger.error(error_message) + for req in reqs: + prepare_abort( + req, + error_message, + status_code=HTTPStatus.INTERNAL_SERVER_ERROR, + ) + if self.scheduler.enable_metrics: + self.scheduler.metrics_collector.increment_bootstrap_failed_reqs() + self.scheduler.stream_output([req], req.return_logprob) + del self._bootstrap_addr_retry_count[bootstrap_addr] + else: + remaining.extend(reqs) continue + if bootstrap_addr in self._bootstrap_addr_retry_count: + del self._bootstrap_addr_retry_count[bootstrap_addr] + need_query = [] for req in reqs: # NOTE: we need resolve it again because we may ensure the parallel info here