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

Co-authored-by: Shangming Cai <csmthu@gmail.com>

python/sglang/srt/disaggregation/common/conn.py

@@ -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

@@ -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

@@ -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

@@ -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(
                         [

test/registered/unit/disaggregation/test_nixl_backend_basic.py

@@ -612,6 +612,7 @@ class TestNixlReceiverPoll(CustomTestCase):
         receiver.init_time = None
         receiver.conclude_state = None
         receiver.abort_notified = False
+        receiver._connection_pool_entries = {}
         return receiver, mgr
 
     def test_returns_existing_conclude_state_without_polling_manager(self):

test/registered/unit/disaggregation/test_receiver_connection_pool.py

@@ -0,0 +1,143 @@
+"""Unit tests for srt/disaggregation/common/conn — receiver connection_pool invalidation."""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+import threading
+import unittest
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+from sglang.srt.disaggregation.base.conn import KVPoll
+from sglang.srt.disaggregation.common.conn import CommonKVReceiver
+from sglang.test.test_utils import CustomTestCase
+
+
+class _ConcreteReceiver(CommonKVReceiver):
+    def poll(self) -> KVPoll:
+        raise NotImplementedError
+
+    def failure_exception(self):
+        raise NotImplementedError
+
+
+def _receiver(connection_pool, entries):
+    receiver = object.__new__(_ConcreteReceiver)
+    receiver.kv_mgr = SimpleNamespace(
+        connection_pool=connection_pool,
+        connection_lock=threading.Lock(),
+    )
+    receiver._connection_pool_entries = entries
+    return receiver
+
+
+class _FetchingReceiver(_ConcreteReceiver):
+    def _get_bootstrap_info_from_server(
+        self, prefill_dp_rank, prefill_cp_rank, target_tp_rank, target_pp_rank
+    ):
+        self.fetch_count += 1
+        return {"rank_ip": "10.0.0.1", "rank_port": 2001, "pp_rank": target_pp_rank}
+
+    def _register_kv_args(self):
+        return True
+
+
+def _fetching_receiver(connection_pool):
+    receiver = object.__new__(_FetchingReceiver)
+    receiver.kv_mgr = SimpleNamespace(
+        connection_pool=connection_pool,
+        connection_lock=threading.Lock(),
+        is_mla_backend=False,
+    )
+    receiver.bootstrap_addr = "prefill:8998"
+    receiver.bootstrap_room = 1
+    receiver.prefill_dp_rank = 0
+    receiver.target_cp_ranks = [0]
+    receiver.target_tp_rank = 0
+    receiver.target_tp_ranks = [0]
+    receiver.target_pp_ranks = [0]
+    receiver._connection_pool_entries = {}
+    receiver.fetch_count = 0
+    return receiver
+
+
+class TestReceiverConnectionPool(CustomTestCase):
+    def test_invalidate_removes_matching_generation(self):
+        stale = [
+            {"rank_ip": "10.0.0.1", "rank_port": 1001},
+            {"rank_ip": "10.0.0.1", "rank_port": 1002},
+        ]
+        retained = [{"rank_ip": "10.0.0.2", "rank_port": 2001}]
+        receiver = _receiver(
+            {"stale": stale, "retained": retained},
+            {"stale": stale},
+        )
+
+        receiver.invalidate_cached_bootstrap_infos()
+
+        self.assertEqual(receiver.kv_mgr.connection_pool, {"retained": retained})
+        self.assertEqual(receiver._connection_pool_entries, {})
+
+    def test_invalidate_preserves_concurrent_replacement_generation(self):
+        stale = [{"rank_ip": "10.0.0.1", "rank_port": 1001}]
+        replacement = [{"rank_ip": "10.0.0.1", "rank_port": 2001}]
+        receiver = _receiver(
+            {"key": replacement},
+            {"key": stale},
+        )
+
+        receiver.invalidate_cached_bootstrap_infos()
+
+        self.assertEqual(receiver.kv_mgr.connection_pool, {"key": replacement})
+
+    def test_invalidate_removes_all_matching_cp_entries(self):
+        stale_cp0 = [{"rank_ip": "10.0.0.1", "rank_port": 1001}]
+        stale_cp1 = [{"rank_ip": "10.0.0.1", "rank_port": 1002}]
+        receiver = _receiver(
+            {"cp0": stale_cp0, "cp1": stale_cp1},
+            {"cp0": stale_cp0, "cp1": stale_cp1},
+        )
+
+        receiver.invalidate_cached_bootstrap_infos()
+
+        self.assertEqual(receiver.kv_mgr.connection_pool, {})
+
+    def test_next_receiver_refetches_after_invalidation(self):
+        stale = [{"rank_ip": "10.0.0.1", "rank_port": 1001}]
+        connection_pool = {"prefill:8998_0_0_0": stale}
+        stale_receiver = _receiver(
+            connection_pool,
+            {"prefill:8998_0_0_0": stale},
+        )
+        stale_receiver.invalidate_cached_bootstrap_infos()
+
+        receiver = _fetching_receiver(connection_pool)
+        receiver._setup_bootstrap_infos()
+
+        self.assertEqual(receiver.fetch_count, 1)
+        self.assertEqual(receiver.bootstrap_infos[0]["rank_port"], 2001)
+        self.assertIs(
+            connection_pool["prefill:8998_0_0_0"],
+            receiver._connection_pool_entries["prefill:8998_0_0_0"],
+        )
+
+    @patch("sglang.srt.disaggregation.common.conn.time.time", return_value=3.0)
+    def test_waiting_timeout_invalidates_cached_generation(self, _mock_time):
+        stale = [{"rank_ip": "10.0.0.1", "rank_port": 1001}]
+        receiver = _receiver({"key": stale}, {"key": stale})
+        receiver.bootstrap_room = 1
+        receiver.bootstrap_infos = stale
+        receiver.init_time = 1.0
+        receiver.abort_notified = True
+        receiver.kv_mgr.waiting_timeout = 1.0
+        receiver.kv_mgr.record_failure = Mock()
+        receiver.kv_mgr.update_status = Mock()
+
+        self.assertEqual(receiver._check_waiting_timeout(), KVPoll.Failed)
+        self.assertEqual(receiver.kv_mgr.connection_pool, {})
+
+
+if __name__ == "__main__":
+    unittest.main()