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

[Router] Publish cache-aware load state (#38139)

Browse Source
This commit is contained in:
Vincent Gao
2026-09-06 15:22:25 +08:00
committed by GitHub
parent 6cee9285a3
commit ae54ccb25d
2 changed files with 53 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/managers/test_loadstat_wire.py
+29 -11
View File
@@ -1,11 +1,13 @@
"""Wire contract and port/rank gating for the LoadStat load snapshot.
Locks the msgpack array shape the sgl-router `cache_aware_zmq` policy will
decode positionally (that consumer lands with the router PR; it is not yet
in this tree, so this pins only the Python side):
Locks the msgpack array shape the sgl-router engine load subscriber in #38108
will decode positionally (that consumer is not yet in this tree, so this pins
only the Python side):
["LoadStat", num_running_reqs, num_waiting_reqs, num_tokens,
max_total_num_tokens, attn_dp_rank]
max_total_num_tokens, attn_dp_rank,
num_waiting_uncached_tokens, num_total_tokens, max_running_requests,
total_prefill_uncached_tokens, total_prefill_busy_us]
carried as the payload of a three-frame message ``[b"load", BE-i64 seq,
payload]``. A field reorder or rename is a silent cross-language break, so
@@ -46,7 +48,7 @@ class TestLoadStatWire(CustomTestCase):
attn_dp_rank=2,
)
)
self.assertEqual(raw.hex(), "96a84c6f6164537461740703cd0400cd200002")
self.assertEqual(raw.hex(), "9ba84c6f6164537461740703cd0400cd2000020000000000")
def test_loadstat_msgpack_array_shape(self):
raw = msgspec.msgpack.Encoder().encode(
@@ -59,10 +61,10 @@ class TestLoadStatWire(CustomTestCase):
)
)
# tag=True + array_like → [tag, *fields] in declaration order; the
# router reads the tag + four counts and ignores the trailing field.
# router reads the tag and four base fields and ignores the suffix.
self.assertEqual(
msgspec.msgpack.Decoder().decode(raw),
["LoadStat", 7, 3, 1024, 8192, 2],
["LoadStat", 7, 3, 1024, 8192, 2, 0, 0, 0, 0, 0],
)
def test_loadstat_tag_is_class_name(self):
@@ -78,8 +80,9 @@ class TestLoadStatWire(CustomTestCase):
)
decoded = msgspec.msgpack.Decoder().decode(raw)
# LoadStat sets no omit_defaults, so the trailing field is always
# emitted (null when unset); a decoder must tolerate it.
self.assertEqual(decoded, ["LoadStat", 0, 0, 0, 0, None])
# emitted (null when unset); a decoder must tolerate it. The new suffix
# fields are always emitted as well.
self.assertEqual(decoded, ["LoadStat", 0, 0, 0, 0, None, 0, 0, 0, 0, 0])
ZMQ_ENDPOINT = '{"publisher": "zmq", "endpoint": "tcp://*:5557"}'
@@ -336,6 +339,11 @@ class TestLoadPublisherGating(CustomTestCase):
num_waiting_reqs=2,
num_used_tokens=3,
max_total_num_tokens=4,
num_waiting_uncached_tokens=5,
num_total_tokens=6,
max_running_requests=7,
total_prefill_uncached_tokens=8,
total_prefill_busy_us=9,
)
)
@@ -356,6 +364,11 @@ class TestLoadPublisherGating(CustomTestCase):
num_waiting_reqs=2,
num_used_tokens=3,
max_total_num_tokens=4,
num_waiting_uncached_tokens=5,
num_total_tokens=6,
max_running_requests=7,
total_prefill_uncached_tokens=8,
total_prefill_busy_us=9,
)
pub.publish_load_stat(provider, force=True, snapshot=snap)
provider.assert_not_called()
@@ -372,7 +385,7 @@ class TestLoadPublisherGating(CustomTestCase):
self.assertEqual(seq, (0).to_bytes(8, "big"))
self.assertEqual(
msgspec.msgpack.Decoder().decode(payload),
["LoadStat", 1, 2, 3, 4, 0],
["LoadStat", 1, 2, 3, 4, 0, 5, 6, 7, 8, 9],
)
def test_unchanged_stat_is_deduped_to_the_heartbeat(self):
@@ -475,6 +488,11 @@ class TestLoadStatIntegration(CustomTestCase):
num_waiting_reqs=3,
num_used_tokens=1024,
max_total_num_tokens=8192,
num_waiting_uncached_tokens=512,
num_total_tokens=4096,
max_running_requests=64,
total_prefill_uncached_tokens=20_000,
total_prefill_busy_us=2_000_000,
)
# PUB/SUB drops messages sent before the subscription propagates, so
# re-publish until one lands (heartbeat reset each pass).
@@ -492,7 +510,7 @@ class TestLoadStatIntegration(CustomTestCase):
self.assertEqual(len(seq), 8)
self.assertEqual(
msgspec.msgpack.Decoder().decode(payload),
["LoadStat", 7, 3, 1024, 8192, 0],
["LoadStat", 7, 3, 1024, 8192, 0, 512, 4096, 64, 20_000, 2_000_000],
)