diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 7b93c26c1..82832a429 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -183,13 +183,26 @@ class CommonKVManager(BaseKVManager): with self.failure_lock: self.failure_records[bootstrap_room] = failure_reason - def ensure_parallel_info(self, bootstrap_addr: str) -> bool: + def ensure_parallel_info( + 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 = self._fetch_prefill_server_info(bootstrap_addr) + 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) if info is None: return False @@ -414,12 +427,10 @@ class CommonKVReceiver(BaseKVReceiver): self.kv_mgr = mgr self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping) - if self.bootstrap_addr not in self.kv_mgr.prefill_info_table: + if not self.kv_mgr.ensure_parallel_info(self.bootstrap_addr): 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 36e99b28b..9d55ba027 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -265,7 +265,6 @@ 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: @@ -514,33 +513,19 @@ 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): - # 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) + 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) 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