From 12408ec668bedb7cf427f58e401491d46f158ae8 Mon Sep 17 00:00:00 2001 From: zhaochen zhang Date: Fri, 15 May 2026 16:38:51 +0800 Subject: [PATCH] =?UTF-8?q?[Disagg]=20Add=20retry=20with=20exponential=20b?= =?UTF-8?q?ackoff=20for=20prefill=20bootstrap=20reg=E2=80=A6=20(#25125)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../sglang/srt/disaggregation/common/conn.py | 37 ++- .../test_register_to_bootstrap.py | 226 ++++++++++++++++++ 2 files changed, 252 insertions(+), 11 deletions(-) create mode 100644 test/registered/unit/disaggregation/test_register_to_bootstrap.py diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index f2a00d59a..596b58303 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -374,7 +374,7 @@ class CommonKVManager(BaseKVManager): return synced_port def register_to_bootstrap(self): - """Register prefill server info to bootstrap server via HTTP POST.""" + """Register prefill server info to bootstrap server via HTTP PUT.""" if self.dist_init_addr: # Multi-node case: bootstrap server's host is dist_init_addr host = NetworkAddress.parse(self.dist_init_addr).resolved().host @@ -402,18 +402,33 @@ class CommonKVManager(BaseKVManager): "load_balance_method": self.server_args.load_balance_method, } - try: - response = requests.put(url, json=payload, timeout=5) - if response.status_code == 200: - logger.debug("Prefill successfully registered to bootstrap server.") - else: - logger.error( - f"Prefill instance failed to connect to bootstrap server: {response.status_code}, {response.text}" + max_retries, initial_delay, max_delay = 5, 1.0, 30.0 + for attempt in range(max_retries): + try: + response = requests.put(url, json=payload, timeout=5) + if response.status_code == 200: + logger.debug("Prefill successfully registered to bootstrap server.") + return + logger.warning( + f"Prefill register attempt {attempt + 1}/{max_retries} failed: status {response.status_code}" ) - except Exception as e: - logger.error( - f"Prefill instance failed to register to bootstrap server: {e}" + except Exception as e: + # Walk to root cause to skip misleading urllib3 wrapper messages + cause = e + while cause.__cause__ is not None: + cause = cause.__cause__ + logger.warning( + f"Prefill register attempt {attempt + 1}/{max_retries} failed: {cause}" + ) + if attempt == max_retries - 1: + break + delay = min(initial_delay * (2**attempt), max_delay) * ( + 0.75 + 0.25 * (time.monotonic() % 1) ) + time.sleep(delay) + logger.error( + f"Prefill instance failed to register to bootstrap server after {max_retries} retries" + ) @cache def _connect(self, endpoint: str, is_ipv6: bool = False): diff --git a/test/registered/unit/disaggregation/test_register_to_bootstrap.py b/test/registered/unit/disaggregation/test_register_to_bootstrap.py new file mode 100644 index 000000000..b16574812 --- /dev/null +++ b/test/registered/unit/disaggregation/test_register_to_bootstrap.py @@ -0,0 +1,226 @@ +"""Unit tests for srt/disaggregation/common/conn — register_to_bootstrap retry logic.""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="stage-a-test-cpu") + +import unittest +from unittest.mock import MagicMock, call, patch + +from sglang.test.test_utils import CustomTestCase + + +class TestRegisterToBootstrap(CustomTestCase): + """Tests for CommonKVManager.register_to_bootstrap retry/backoff behavior.""" + + @patch("sglang.srt.disaggregation.common.conn.time") + @patch("sglang.srt.disaggregation.common.conn.requests.put") + def test_succeeds_on_first_attempt(self, mock_put, mock_time): + mock_time.monotonic.return_value = 0.0 + mock_response = MagicMock() + mock_response.status_code = 200 + mock_put.return_value = mock_response + + mgr = self._make_manager() + mgr.register_to_bootstrap() + + mock_put.assert_called_once() + mock_time.sleep.assert_not_called() + + @patch("sglang.srt.disaggregation.common.conn.time") + @patch("sglang.srt.disaggregation.common.conn.requests.put") + def test_succeeds_after_retries(self, mock_put, mock_time): + mock_time.monotonic.return_value = 0.0 + fail_resp = MagicMock() + fail_resp.status_code = 503 + success_resp = MagicMock() + success_resp.status_code = 200 + mock_put.side_effect = [fail_resp, fail_resp, success_resp] + + mgr = self._make_manager() + mgr.register_to_bootstrap() + + self.assertEqual(mock_put.call_count, 3) + self.assertEqual(mock_time.sleep.call_count, 2) + + @patch("sglang.srt.disaggregation.common.conn.time") + @patch("sglang.srt.disaggregation.common.conn.requests.put") + def test_all_retries_exhausted(self, mock_put, mock_time): + mock_time.monotonic.return_value = 0.0 + fail_resp = MagicMock() + fail_resp.status_code = 503 + mock_put.return_value = fail_resp + + mgr = self._make_manager() + mgr.register_to_bootstrap() + + self.assertEqual(mock_put.call_count, 5) + # Sleep is only called between attempts, not after the final failure + self.assertEqual(mock_time.sleep.call_count, 4) + + @patch("sglang.srt.disaggregation.common.conn.time") + @patch("sglang.srt.disaggregation.common.conn.requests.put") + def test_exception_with_nested_cause(self, mock_put, mock_time): + mock_time.monotonic.return_value = 0.0 + + root_exc = ConnectionRefusedError("connection refused") + inner_exc = OSError("os error") + inner_exc.__cause__ = root_exc + outer_exc = Exception("wrapped") + outer_exc.__cause__ = inner_exc + + success_resp = MagicMock() + success_resp.status_code = 200 + mock_put.side_effect = [outer_exc, success_resp] + + mgr = self._make_manager() + mgr.register_to_bootstrap() + + self.assertEqual(mock_put.call_count, 2) + + @patch("sglang.srt.disaggregation.common.conn.time") + @patch("sglang.srt.disaggregation.common.conn.requests.put") + def test_exception_with_no_cause(self, mock_put, mock_time): + mock_time.monotonic.return_value = 0.0 + + exc = ConnectionError("plain connection error") + exc.__cause__ = None + + success_resp = MagicMock() + success_resp.status_code = 200 + mock_put.side_effect = [exc, success_resp] + + mgr = self._make_manager() + mgr.register_to_bootstrap() + + self.assertEqual(mock_put.call_count, 2) + + @patch("sglang.srt.disaggregation.common.conn.time") + @patch("sglang.srt.disaggregation.common.conn.requests.put") + def test_backoff_delay_exponential(self, mock_put, mock_time): + mock_time.monotonic.return_value = 0.0 + fail_resp = MagicMock() + fail_resp.status_code = 503 + mock_put.return_value = fail_resp + + mgr = self._make_manager() + mgr.register_to_bootstrap() + + # With monotonic() = 0.0, jitter factor = 0.75 + 0.25 * (0.0 % 1) = 0.75 + # delay = min(1.0 * 2^attempt, 30.0) * 0.75 + # Sleep happens only between attempts (attempt 0..3), not after the final failure + expected_calls = [call(0.75), call(1.5), call(3.0), call(6.0)] + self.assertEqual(mock_time.sleep.call_args_list, expected_calls) + + @patch("sglang.srt.disaggregation.common.conn.time") + @patch("sglang.srt.disaggregation.common.conn.requests.put") + def test_jitter_never_exceeds_max_delay(self, mock_put, mock_time): + """Guard against operator-precedence regressions in the jitter factor. + + The jitter factor must stay in [0.75, 1.0), so a delay capped at + max_delay must never exceed max_delay after applying jitter. + """ + # monotonic() returns a value whose fractional part is close to 1. + # If the parentheses around `time.monotonic() % 1` were dropped, the + # jitter factor could grow up to ~1.75 and blow past max_delay. + mock_time.monotonic.return_value = 999.9999 + fail_resp = MagicMock() + fail_resp.status_code = 503 + mock_put.return_value = fail_resp + + mgr = self._make_manager() + mgr.register_to_bootstrap() + + max_delay = 30.0 + for sleep_call in mock_time.sleep.call_args_list: + actual_delay = sleep_call[0][0] + self.assertLess(actual_delay, max_delay) + self.assertGreaterEqual(actual_delay, 0.75) + + @patch("sglang.srt.disaggregation.common.conn.time") + @patch("sglang.srt.disaggregation.common.conn.requests.put") + def test_payload_contains_required_fields(self, mock_put, mock_time): + mock_time.monotonic.return_value = 0.0 + success_resp = MagicMock() + success_resp.status_code = 200 + mock_put.return_value = success_resp + + mgr = self._make_manager() + mgr.register_to_bootstrap() + + call_kwargs = mock_put.call_args + payload = call_kwargs[1]["json"] + required_fields = [ + "attn_tp_size", + "attn_tp_rank", + "attn_cp_size", + "attn_cp_rank", + "attn_dp_size", + "attn_dp_rank", + "pp_size", + "pp_rank", + "system_dp_size", + "system_dp_rank", + "rank_ip", + "rank_port", + "page_size", + "kv_cache_dtype", + ] + for field in required_fields: + self.assertIn(field, payload) + + @patch("sglang.srt.disaggregation.common.conn.time") + @patch("sglang.srt.disaggregation.common.conn.requests.put") + def test_url_with_dist_init_addr(self, mock_put, mock_time): + mock_time.monotonic.return_value = 0.0 + success_resp = MagicMock() + success_resp.status_code = 200 + mock_put.return_value = success_resp + + mgr = self._make_manager(dist_init_addr="10.0.0.1:12345") + mgr.register_to_bootstrap() + + url_used = mock_put.call_args[0][0] + self.assertIn("10.0.0.1", url_used) + + def _make_manager(self, dist_init_addr=None): + """Create a lightweight mock manager that has the attributes needed + by register_to_bootstrap, without going through CommonKVManager.__init__ + (which requires zmq, ServerArgs model resolution, etc.).""" + from sglang.srt.disaggregation.common.conn import CommonKVManager + + mgr = MagicMock(spec=CommonKVManager) + # Bind the real method to the mock + mgr.register_to_bootstrap = CommonKVManager.register_to_bootstrap.__get__( + mgr, CommonKVManager + ) + + # Set attributes that register_to_bootstrap reads + mgr.dist_init_addr = dist_init_addr + mgr.bootstrap_host = "127.0.0.1" + mgr.bootstrap_port = 8765 + mgr.attn_tp_size = 1 + mgr.attn_tp_rank = 0 + mgr.attn_cp_size = 1 + mgr.attn_cp_rank = 0 + mgr.attn_dp_size = 1 + mgr.attn_dp_rank = 0 + mgr.pp_size = 1 + mgr.pp_rank = 0 + mgr.system_dp_size = 1 + mgr.system_dp_rank = 0 + mgr.local_ip = "127.0.0.1" + mgr.rank_port = 12345 + + mgr.kv_args = MagicMock() + mgr.kv_args.page_size = 16 + + mgr.server_args = MagicMock() + mgr.server_args.kv_cache_dtype = "auto" + mgr.server_args.load_balance_method = "follow_bootstrap_room" + + return mgr + + +if __name__ == "__main__": + unittest.main()