diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index ab924627c..0688fb193 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -184,7 +184,7 @@ class CommonKVManager(BaseKVManager): self.failure_records[bootstrap_room] = failure_reason def ensure_parallel_info( - self, bootstrap_addr: str, max_retries: int = 20, retry_interval: float = 1.0 + 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). diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 958cff1cd..9d55ba027 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -24,7 +24,7 @@ import logging from collections import deque from dataclasses import dataclass from http import HTTPStatus -from typing import TYPE_CHECKING, List, Optional, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch from torch.distributed import ProcessGroup @@ -498,42 +498,58 @@ class DecodePreallocQueue: if not self.pending_reqs: return - bootstrap_addr = f"{self.pending_reqs[0].bootstrap_host}:{self.pending_reqs[0].bootstrap_port}" - - # If a request is following the bootstrap room, - # we need get the prefill info before resolving the prefill_dp_ranks - # 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): - return + # Group pending requests by bootstrap_addr + addr_to_reqs: Dict[str, List[Req]] = {} + for req in self.pending_reqs: + addr = f"{req.bootstrap_host}:{req.bootstrap_port}" + addr_to_reqs.setdefault(addr, []).append(req) resolved = [] - need_query = [] - for req in self.pending_reqs: - # NOTE: we need resolve it again because we may ensure the parallel info here - prefill_dp_rank = self._resolve_prefill_dp_rank(req) - if prefill_dp_rank is not None: - resolved.append((req, prefill_dp_rank)) - else: - need_query.append(req) + remaining = [] - if need_query: - from sglang.srt.disaggregation.common.conn import CommonKVReceiver + for bootstrap_addr, reqs in addr_to_reqs.items(): + # If a request is following the bootstrap room, + # we need get the prefill info before resolving the prefill_dp_ranks + # 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) + continue - rooms = [req.bootstrap_room for req in need_query] - room_to_rank = CommonKVReceiver.query_prefill_dp_ranks( - bootstrap_addr, rooms - ) - remaining = [] - for req in need_query: - room_key = str(req.bootstrap_room) - if room_key in room_to_rank: - resolved.append((req, int(room_to_rank[room_key]))) + need_query = [] + for req in reqs: + # NOTE: we need resolve it again because we may ensure the parallel info here + prefill_dp_rank = self._resolve_prefill_dp_rank(req) + if prefill_dp_rank is not None: + resolved.append((req, prefill_dp_rank)) else: - remaining.append(req) - self.pending_reqs = remaining - else: - self.pending_reqs = [] + need_query.append(req) + + if need_query: + from sglang.srt.disaggregation.common.conn import CommonKVReceiver + + rooms = [req.bootstrap_room for req in need_query] + room_to_rank = CommonKVReceiver.query_prefill_dp_ranks( + bootstrap_addr, rooms + ) + for req in need_query: + room_key = str(req.bootstrap_room) + if room_key in room_to_rank: + resolved.append((req, int(room_to_rank[room_key]))) + else: + remaining.append(req) + + self.pending_reqs = remaining for req, prefill_dp_rank in resolved: self._create_receiver_and_enqueue(req, prefill_dp_rank)