[Disagg] Add retry with exponential backoff for prefill bootstrap reg… (#25125)
This commit is contained in:
zhaochen zhang
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user