minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix(disagg): refresh stale prefill bootstrap metadata (#36029)

Browse Source 
Co-authored-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Kevin Flansburg
2026-08-26 03:29:20 +08:00
committed by GitHub
co-authored by Shangming Cai
parent 99c02d71b1
commit 7ddf92d5f4
6 changed files with 180 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/common/conn.py
+27 -3
View File
@@ -1357,6 +1357,7 @@ class CommonKVReceiver(BaseKVReceiver):
self.require_staging: bool = False
self.init_time: Optional[float] = None
self.abort_notified: bool = False
self._connection_pool_entries: Dict[str, List[Dict]] = {}
self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room)
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping)
@@ -1403,7 +1404,10 @@ class CommonKVReceiver(BaseKVReceiver):
for target_cp_rank in self.target_cp_ranks:
bootstrap_key = f"{self.bootstrap_addr}_{self.prefill_dp_rank}_{target_cp_rank}_{self.target_tp_rank}"
if bootstrap_key not in self.kv_mgr.connection_pool:
with self.kv_mgr.connection_lock:
cached_bootstrap_infos = self.kv_mgr.connection_pool.get(bootstrap_key)
if cached_bootstrap_infos is None:
bootstrap_infos = []
for target_tp_rank in self.target_tp_ranks:
# Enable higher PP ranks to be bootstrapped earlier to make PP PD requests bootstrap more robust
@@ -1438,24 +1442,43 @@ class CommonKVReceiver(BaseKVReceiver):
self.bootstrap_room, KVPoll.Failed
)
self.bootstrap_infos = None
self.invalidate_cached_bootstrap_infos()
return
self.bootstrap_infos = bootstrap_infos
self._connection_pool_entries[bootstrap_key] = self.bootstrap_infos
# Register kv_args only once to prefill KVManager according to the info fetched
# from the bootstrap server. Do this before caching in connection_pool so a failed
# registration does not leave a stale entry that later requests would reuse.
if not self._register_kv_args():
self.invalidate_cached_bootstrap_infos()
return
self.kv_mgr.connection_pool[bootstrap_key] = self.bootstrap_infos
with self.kv_mgr.connection_lock:
cached_bootstrap_infos = self.kv_mgr.connection_pool.setdefault(
bootstrap_key, self.bootstrap_infos
)
if cached_bootstrap_infos is not self.bootstrap_infos:
self.bootstrap_infos = cached_bootstrap_infos
else:
self.bootstrap_infos = self.kv_mgr.connection_pool[bootstrap_key]
self.bootstrap_infos = cached_bootstrap_infos
self._connection_pool_entries[bootstrap_key] = self.bootstrap_infos
assert len(self.bootstrap_infos) > 0
all_bootstrap_infos.extend(self.bootstrap_infos)
self.bootstrap_infos = all_bootstrap_infos
def invalidate_cached_bootstrap_infos(self) -> None:
with self.kv_mgr.connection_lock:
for bootstrap_key, bootstrap_infos in self._connection_pool_entries.items():
if self.kv_mgr.connection_pool.get(bootstrap_key) is bootstrap_infos:
del self.kv_mgr.connection_pool[bootstrap_key]
self._connection_pool_entries.clear()
def _get_bootstrap_info_from_server(
self, prefill_dp_rank, prefill_cp_rank, target_tp_rank, target_pp_rank
):
@@ -1566,6 +1589,7 @@ class CommonKVReceiver(BaseKVReceiver):
f"in KVPoll.WaitingForInput",
)
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
self.invalidate_cached_bootstrap_infos()
if (
not self.abort_notified
and hasattr(self, "bootstrap_infos")
python/sglang/srt/disaggregation/mooncake/conn.py
+3 -2
View File
@@ -2419,8 +2419,8 @@ class MooncakeKVReceiver(MooncakeFailureExceptionMixin, CommonKVReceiver):
packed_staging_base_ptr = b""
staging_total_size_str = b""
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
try:
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
with lock:
sock.send_multipart(
[
@@ -2480,9 +2480,9 @@ class MooncakeKVReceiver(MooncakeFailureExceptionMixin, CommonKVReceiver):
)
for bootstrap_info in self.bootstrap_infos:
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
is_dummy = bootstrap_info["is_dummy"]
try:
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
with lock:
sock.send_multipart(
[
@@ -2507,6 +2507,7 @@ class MooncakeKVReceiver(MooncakeFailureExceptionMixin, CommonKVReceiver):
]
)
except zmq.ZMQError:
self.invalidate_cached_bootstrap_infos()
self.kv_mgr.record_failure(
self.bootstrap_room,
f"send_metadata to prefill {bootstrap_info.get('rank_ip')}:{bootstrap_info.get('rank_port')} failed",
python/sglang/srt/disaggregation/mori/conn.py
+3 -2
View File
@@ -1683,8 +1683,8 @@ class MoriKVReceiver(CommonKVReceiver):
)
for bootstrap_info in self.bootstrap_infos:
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
try:
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
with lock:
sock.send_multipart(
[
@@ -1737,13 +1737,13 @@ class MoriKVReceiver(CommonKVReceiver):
)
for bootstrap_info in self.bootstrap_infos:
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
is_dummy = bootstrap_info.get("is_dummy", False)
if not is_dummy and normalized_state is not None:
state_bytes = _pack_state_indices(normalized_state)
else:
state_bytes = b""
try:
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
with lock:
sock.send_multipart(
[
@@ -1760,6 +1760,7 @@ class MoriKVReceiver(CommonKVReceiver):
]
)
except zmq.ZMQError:
self.invalidate_cached_bootstrap_infos()
self.kv_mgr.record_failure(
self.bootstrap_room,
f"send_metadata to prefill {bootstrap_info.get('rank_ip')}:{bootstrap_info.get('rank_port')} failed",
python/sglang/srt/disaggregation/nixl/conn.py
+3 -2
View File
@@ -2874,7 +2874,6 @@ class NixlKVReceiver(CommonKVReceiver):
logger.debug(
f"Fetched bootstrap info: {bootstrap_info} for engine rank: {self.kv_mgr.kv_args.engine_rank}"
)
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
is_dummy = bootstrap_info["is_dummy"]
logger.debug(
f"Sending to prefill server with bootstrap room {self.bootstrap_room} {is_dummy=}"
@@ -2887,6 +2886,7 @@ class NixlKVReceiver(CommonKVReceiver):
else b""
)
try:
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
with lock:
sock.send_multipart(
[
@@ -2904,6 +2904,7 @@ class NixlKVReceiver(CommonKVReceiver):
]
)
except zmq.ZMQError:
self.invalidate_cached_bootstrap_infos()
self.kv_mgr.record_failure(
self.bootstrap_room,
f"send_metadata to prefill {bootstrap_info.get('rank_ip')}:{bootstrap_info.get('rank_port')} failed",
@@ -3002,8 +3003,8 @@ class NixlKVReceiver(CommonKVReceiver):
dst_kv_item_len = 0
dst_num_slots = 0
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
try:
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
with lock:
sock.send_multipart(
[