[PD] Make pending reqs resolving more robust (#20505)

Signed-off-by: Shangming Cai <csmthu@gmail.com>

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(
+    def ensure_parallel_info(self, bootstrap_addr: str) -> bool:
-        self, bootstrap_addr: str, max_retries: int = 5, retry_interval: float = 1.0
-    ) -> 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
+        info = self._fetch_prefill_server_info(bootstrap_addr)
-        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)
         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

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}"
+                # Increment retry count for this bootstrap_addr
-                logger.error(error_message)
+                retry_count = (
-                for req in reqs:
+                    self._bootstrap_addr_retry_count.get(bootstrap_addr, 0) + 1
-                    prepare_abort(
+                )
-                        req,
+                self._bootstrap_addr_retry_count[bootstrap_addr] = retry_count
-                        error_message,
+
-                        status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+                if retry_count >= 5:
-                    )
+                    # Max retries reached, fail the requests
-                    if self.scheduler.enable_metrics:
+                    error_message = f"Could not fetch prefill parallel info from bootstrap server {bootstrap_addr} after {retry_count} retries"
-                        self.scheduler.metrics_collector.increment_bootstrap_failed_reqs()
+                    logger.error(error_message)
-                    self.scheduler.stream_output([req], req.return_logprob)
+                    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