diff --git a/python/sglang/srt/managers/scheduler_components/load_publisher.py b/python/sglang/srt/managers/scheduler_components/load_publisher.py index b5bada590..f30ca0974 100644 --- a/python/sglang/srt/managers/scheduler_components/load_publisher.py +++ b/python/sglang/srt/managers/scheduler_components/load_publisher.py @@ -1,10 +1,9 @@ """Per-scheduler load reporting for load-aware routers. Each scheduler publishes a periodic `LoadStat` gauge on its own ZMQ PUB -socket so out-of-process load-aware routers (e.g. sgl-router's -`cache_aware_zmq` policy) can route on real queue depth instead of a -router-side in-flight counter. The in-deployment counterpart lives in -`sglang.srt.managers.load_snapshot` (SHM / PUSH to node 0), which a router +socket so out-of-process load-aware routers can route on real queue depth +instead of a router-side in-flight counter. The in-deployment counterpart +lives in `sglang.srt.managers.load_snapshot` (SHM / PUSH to node 0), which a router that only knows the worker URL cannot subscribe to; the port is instead advertised via `/server_info` (`runtime_context.describe_kv_events_publisher`). The payload is a compact tagged subset of `LoadSnapshot` so the wire @@ -80,9 +79,12 @@ class LoadStat( """Per-scheduler runtime load snapshot. Wire shape (tag + array_like): ``["LoadStat", num_running_reqs, - num_waiting_reqs, num_tokens, max_total_num_tokens, attn_dp_rank]``. The - router reads the four counts; array_like always emits the trailing field - (null when unset), so a decoder must tolerate it. + num_waiting_reqs, num_tokens, 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]``. + + The first six positions are the #34608-compatible prefix. Current routers + ignore appended fields, so extending the payload preserves compatibility. """ num_running_reqs: int @@ -92,6 +94,11 @@ class LoadStat( # attn_dp_rank under DP attention, else the plain dp_rank; informational # only (the router keys by socket rank). Name follows EventBatch's. attn_dp_rank: Optional[int] = None + num_waiting_uncached_tokens: int = 0 # Uncached prompt tokens awaiting prefill + num_total_tokens: int = 0 # KV tokens in use plus queued request tokens + max_running_requests: int = 0 # Scheduler running-request limit + total_prefill_uncached_tokens: int = 0 # Cumulative uncached prefill tokens + total_prefill_busy_us: int = 0 # Cumulative prefill step time in microseconds def _open_pub_socket(endpoint: str) -> zmq.Socket: @@ -228,6 +235,11 @@ class SchedulerLoadPublisher: load.num_waiting_reqs, load.num_used_tokens, load.max_total_num_tokens, + load.num_waiting_uncached_tokens, + load.num_total_tokens, + load.max_running_requests, + load.total_prefill_uncached_tokens, + load.total_prefill_busy_us, ) if ( counts == self._last_counts @@ -241,6 +253,11 @@ class SchedulerLoadPublisher: num_tokens=counts[2], max_total_num_tokens=counts[3], attn_dp_rank=self._rank, + num_waiting_uncached_tokens=counts[4], + num_total_tokens=counts[5], + max_running_requests=counts[6], + total_prefill_uncached_tokens=counts[7], + total_prefill_busy_us=counts[8], ) ) seq = next(self._seq).to_bytes(8, "big") diff --git a/test/registered/unit/managers/test_loadstat_wire.py b/test/registered/unit/managers/test_loadstat_wire.py index 56fc2692a..3211ac810 100644 --- a/test/registered/unit/managers/test_loadstat_wire.py +++ b/test/registered/unit/managers/test_loadstat_wire.py @@ -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], )