From 7ddf92d5f481e79abd6db4f6c28dd27641fd2a7e Mon Sep 17 00:00:00 2001 From: Kevin Flansburg Date: Tue, 25 Aug 2026 12:29:20 -0700 Subject: [PATCH] fix(disagg): refresh stale prefill bootstrap metadata (#36029) Co-authored-by: Shangming Cai --- .../sglang/srt/disaggregation/common/conn.py | 30 +++- .../srt/disaggregation/mooncake/conn.py | 5 +- python/sglang/srt/disaggregation/mori/conn.py | 5 +- python/sglang/srt/disaggregation/nixl/conn.py | 5 +- .../disaggregation/test_nixl_backend_basic.py | 1 + .../test_receiver_connection_pool.py | 143 ++++++++++++++++++ 6 files changed, 180 insertions(+), 9 deletions(-) create mode 100644 test/registered/unit/disaggregation/test_receiver_connection_pool.py diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 74650d342..4b431809b 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/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") diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index d8e47bf73..1c0955653 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/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", diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 33996f7a7..6772c6a7b 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/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", diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 49314582d..fb0c3460c 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/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( [ diff --git a/test/registered/unit/disaggregation/test_nixl_backend_basic.py b/test/registered/unit/disaggregation/test_nixl_backend_basic.py index 0317832d7..6e6ab2238 100644 --- a/test/registered/unit/disaggregation/test_nixl_backend_basic.py +++ b/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): diff --git a/test/registered/unit/disaggregation/test_receiver_connection_pool.py b/test/registered/unit/disaggregation/test_receiver_connection_pool.py new file mode 100644 index 000000000..b28241d90 --- /dev/null +++ b/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()