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

Publish per-scheduler load on a dedicated socket for load-aware routers (#34608)

Browse Source 
Co-authored-by: Kangyan Zhou <zky314343421@gmail.com>
Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
Shangming Cai
2026-08-27 19:33:18 +08:00
committed by GitHub
co-authored by Kangyan Zhou Zhangheng
parent 94183a8d2b
commit 97ba99067d
10 changed files with 1443 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/entrypoints/test_server_info.py
+197
View File
@@ -112,6 +112,7 @@ class TestServerInfoKvEventsField(CustomTestCase):
kv_events_config=(
'{"publisher": "zmq", "endpoint": "tcp://*:5557", "topic": "kv"}'
),
load_publish_endpoint="auto",
page_size=64,
dp_size=2,
)
@@ -128,10 +129,143 @@ class TestServerInfoKvEventsField(CustomTestCase):
"topic": "kv",
"block_size": 64,
"dp_size": 2,
# Load range packed immediately after the KV range: 5557 + dp_size.
"load_endpoint_port_base": 5559,
"load_topic": "load",
},
)
def test_load_port_skips_an_overlapping_replay_range(self):
# Conventional replay = kv + 1: the load range must be advertised
# past the replay ROUTER range (5558 + dp_size), matching where
# SchedulerLoadPublisher actually binds — both sides resolve it via
# resolve_load_pub_range.
args = ServerArgs(
model_path="dummy",
kv_events_config=(
'{"publisher": "zmq", "endpoint": "tcp://*:5557", '
'"replay_endpoint": "tcp://*:5558"}'
),
load_publish_endpoint="auto",
page_size=64,
dp_size=2,
)
info = _call_server_info_with(args)
self.assertEqual(info["kv_events"]["load_endpoint_port_base"], 5560)
self.assertEqual(info["kv_events"]["load_topic"], "load")
def test_explicit_load_publish_endpoint_moves_the_advertised_base(self):
args = ServerArgs(
model_path="dummy",
kv_events_config='{"publisher": "zmq", "endpoint": "tcp://*:5557"}',
load_publish_endpoint="tcp://*:7000",
page_size=64,
dp_size=2,
)
info = _call_server_info_with(args)
self.assertEqual(info["kv_events"]["load_endpoint_port_base"], 7000)
self.assertEqual(info["kv_events"]["load_topic"], "load")
def test_load_keys_omitted_for_connect_style_kv_endpoint(self):
# A concrete host is connected to rather than bound: the KV-events
# descriptor is still valid (KV events work connect-style), but no
# load range can be bound there, so the load keys must be omitted
# rather than advertising a port nothing listens on.
args = ServerArgs(
model_path="dummy",
kv_events_config=(
'{"publisher": "zmq", "endpoint": "tcp://10.0.0.5:5557"}'
),
load_publish_endpoint="auto",
page_size=64,
dp_size=1,
)
info = _call_server_info_with(args)
self.assertIsNotNone(info["kv_events"])
self.assertEqual(info["kv_events"]["endpoint_host"], "10.0.0.5")
self.assertNotIn("load_endpoint_port_base", info["kv_events"])
self.assertNotIn("load_topic", info["kv_events"])
def test_ipv6_wildcard_endpoint_advertises_bracketed_host_and_load_keys(self):
# "[::]" is a bind-all wildcard: the descriptor must keep the
# brackets (consumers splice tcp://{host}:{port}) and the load
# range resolves right after the KV range.
args = ServerArgs(
model_path="dummy",
kv_events_config='{"publisher": "zmq", "endpoint": "tcp://[::]:5557"}',
load_publish_endpoint="auto",
page_size=64,
dp_size=1,
)
info = _call_server_info_with(args)
self.assertEqual(info["kv_events"]["endpoint_host"], "[::]")
self.assertEqual(info["kv_events"]["endpoint_port_base"], 5557)
self.assertEqual(info["kv_events"]["load_endpoint_port_base"], 5558)
def test_concrete_ipv6_endpoint_advertises_kv_but_not_load(self):
# A concrete IPv6 host works connect-style for KV events (advertised,
# brackets kept) but is not bindable for the load range — and "::"
# appearing inside the address must not be mistaken for a wildcard.
args = ServerArgs(
model_path="dummy",
kv_events_config=(
'{"publisher": "zmq", "endpoint": "tcp://[2001:db8::5]:5557"}'
),
load_publish_endpoint="auto",
page_size=64,
dp_size=1,
)
info = _call_server_info_with(args)
self.assertEqual(info["kv_events"]["endpoint_host"], "[2001:db8::5]")
self.assertNotIn("load_endpoint_port_base", info["kv_events"])
def test_load_keys_omitted_when_explicitly_off(self):
args = ServerArgs(
model_path="dummy",
kv_events_config='{"publisher": "zmq", "endpoint": "tcp://*:5557"}',
load_publish_endpoint="off",
page_size=64,
dp_size=1,
)
info = _call_server_info_with(args)
self.assertIsNotNone(info["kv_events"])
self.assertNotIn("load_endpoint_port_base", info["kv_events"])
self.assertNotIn("load_topic", info["kv_events"])
def test_load_keys_omitted_when_no_load_range_fits(self):
# kv base 65535 leaves no u16 room for a load range: the kv_events
# descriptor must still be served, with only the load keys omitted,
# so routers fall back to their in-flight counter for load.
args = ServerArgs(
model_path="dummy",
kv_events_config='{"publisher": "zmq", "endpoint": "tcp://*:65535"}',
load_publish_endpoint="auto",
page_size=64,
dp_size=1,
)
info = _call_server_info_with(args)
self.assertIsNotNone(info["kv_events"])
self.assertEqual(info["kv_events"]["endpoint_port_base"], 65535)
self.assertNotIn("load_endpoint_port_base", info["kv_events"])
self.assertNotIn("load_topic", info["kv_events"])
def test_kv_events_descriptor_carries_specific_host_and_topic(self):
# No --load-publish-endpoint: KV descriptor served, load keys absent.
# Upgrade safety rests on this default silence, so pin it here.
args = ServerArgs(
model_path="dummy",
kv_events_config=(
@@ -149,6 +283,8 @@ class TestServerInfoKvEventsField(CustomTestCase):
self.assertEqual(info["kv_events"]["topic"], "kv")
self.assertEqual(info["kv_events"]["block_size"], 128)
self.assertEqual(info["kv_events"]["dp_size"], 1)
self.assertNotIn("load_endpoint_port_base", info["kv_events"])
self.assertNotIn("load_topic", info["kv_events"])
# ----- disabled / unconfigured -------------------------------------
@@ -387,5 +523,66 @@ class TestServerInfoExistingFieldsPreserved(CustomTestCase):
json.dumps(info)
class TestLoadPublishEndpointValidation(CustomTestCase):
"""--load-publish-endpoint fails fast at the entrypoint, not silently in a
scheduler subprocess log."""
def test_requires_kv_events_config(self):
args = ServerArgs(model_path="dummy", load_publish_endpoint="tcp://*:6000")
with self.assertRaisesRegex(ValueError, "kv-events"):
args.check_load_publish_args()
def test_rejects_non_bindable_endpoint(self):
args = ServerArgs(
model_path="dummy",
kv_events_config='{"publisher": "zmq", "endpoint": "tcp://*:5557"}',
load_publish_endpoint="tcp://10.0.0.5:6000",
)
with self.assertRaisesRegex(ValueError, "bindable"):
args.check_load_publish_args()
def test_rejects_endpoint_overlapping_the_kv_range(self):
args = ServerArgs(
model_path="dummy",
kv_events_config='{"publisher": "zmq", "endpoint": "tcp://*:5557"}',
dp_size=4,
load_publish_endpoint="tcp://*:5558",
)
with self.assertRaisesRegex(ValueError, "overlaps"):
args.check_load_publish_args()
def test_rejects_null_publisher(self):
# publisher='null' disables KV events, so there is nothing to advertise
# through — accepting the opt-in would silently do nothing.
args = ServerArgs(
model_path="dummy",
kv_events_config='{"publisher": "null"}',
load_publish_endpoint="auto",
)
with self.assertRaisesRegex(ValueError, "null"):
args.check_load_publish_args()
def test_rejects_unparseable_kv_events_config(self):
args = ServerArgs(
model_path="dummy",
kv_events_config="{not json",
load_publish_endpoint="auto",
)
with self.assertRaisesRegex(ValueError, "not parseable"):
args.check_load_publish_args()
def test_off_and_valid_endpoint_pass(self):
for endpoint in (None, "off", "OFF", "auto", "tcp://*:6000"):
with self.subTest(endpoint=endpoint):
args = ServerArgs(
model_path="dummy",
kv_events_config=(
'{"publisher": "zmq", "endpoint": "tcp://*:5557"}'
),
load_publish_endpoint=endpoint,
)
args.check_load_publish_args() # must not raise
if __name__ == "__main__":
unittest.main()