[scheduler] Publish per-scheduler load on a dedicated socket for load-aware routers

[Kangyan-Zhou]

Motivation

Load-aware routers currently infer a worker's load from a router-side in-flight counter. That counter measures what this router has dispatched, not what the engine actually has queued — it misses traffic from other routers and direct clients, and for streaming responses it stays held for the whole response rather than the time the request occupies the scheduler.

The engine already builds exactly the right number. SchedulerLoadInquirer.get_loads() produces a LoadSnapshot every publish cycle, which today feeds /v1/loads and DP-attention dispatch. This exposes that same snapshot to out-of-process consumers, so a router can price workers on real queue depth.

The router-side consumer is #28600.

Modifications

A third writer in the existing load-snapshot family. PubLoadSnapshotWriter joins ShmLoadSnapshotWriter and ZmqLoadSnapshotWriter in managers/load_snapshot.py, behind the same write() / publish_interval / close() interface. It is not a variant of the existing ZMQ writer: that one is PUSH and connects to the single reader process that owns the PULL end, so an additional consumer would load-balance snapshots away from the DP controller rather than receive copies. Routers need fan-out, so this binds PUB.

Collection and publication get one owner. SchedulerLoadInquirer now holds the writers and gains publish(); Scheduler.publish_load_snapshot is deleted and its two call sites delegate. The snapshot is built once per cycle no matter how many writers are due — relevant because get_loads() walks the running batch, the waiting queue and four disaggregation queues.

Port derivation and advertisement share one decision. load_pub_port_base() decides whether a load range exists and where it starts (kv_events_port + dp_size; load rank r uses base + r). Both the writer and /server_info's load_endpoint_port_base route through it, so the engine cannot advertise a range it will not bind — a router subscribing to a port nobody bound waits forever while reporting the worker as an expected publisher.

It declines, and the advertisement is omitted, when:

• there is no kv-events config, or the publisher is null;
• the endpoint is not tcp:// — ipc:// and inproc:// serve KV events fine but have no port to offset;
• the endpoint has a concrete host, which the publisher would connect to rather than bind; nothing listens on the load range, so connecting publishes into a void neither side can detect;
• the range would run past the u16 ceiling.

A router treats a missing advertisement as "this engine does not report load" and falls back to its own signal, so declining is a supported outcome rather than an error.

Wire framing. Subscribers require a three-frame message — [topic, big-endian i64 seq, msgpack payload] — and drop anything else. ZMQ_CONFLATE would be the natural fit for a gauge but keeps only a single frame, which would corrupt that layout; a small send HWM bounds the backlog instead. The topic frame is empty: this socket carries only load, so subscribers subscribe-all.

Failure handling is asymmetric on purpose. Each writer is constructed under its own guard, so a bad router-facing endpoint cannot cost the internal writer. Losing the internal writer is logged as an error — without it /v1/loads omits the rank and the DP controller stops refreshing its budget for it, so the rank drifts out of the dispatch rotation with no other symptom. Publish failures are counted and reported on the first occurrence and every Nth after: the idle path publishes on every scheduler-loop iteration and --sleep-on-idle is off by default, so an unthrottled warning would turn a permanently broken writer into thousands of lines per second.

Notes for reviewers

• Removed override point. Scheduler.publish_load_snapshot was public; forks overriding it will need to move to SchedulerLoadInquirer.publish.
• Port footprint. --kv-events-config now reserves 2 * dp_size consecutive ports from the configured base rather than dp_size. Co-located engines spaced dp_size apart will collide. The argument documentation is not yet updated for this.
• No new server arguments; cadence reuses --load-snapshot-publish-interval.

Accuracy Tests

Not applicable — no changes to model execution or output.

Speed Tests and Profiling

No benchmark run. The one relevant change is a reduction: publishing to N sinks now collects a single snapshot per cycle rather than one per sink, which removes a redundant queue walk on every extend batch.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Tests added in test/registered/unit/managers/test_load_snapshot_backends.py:

• the three-frame layout, against the module's own decoder;
• advertised port equals bound port, in both the accept and decline directions;
• the factory returns both writers, in an order that keeps the internal one from hiding behind the optional one;
• publish() — single collection per cycle, independent per-writer intervals, force semantics, per-writer fault isolation, and the failure throttle;
• the encoder's wire shape is a map keyed by field name, so declaring array_like=True on LoadSnapshot or renaming a field fails here rather than silently breaking consumers.

---

CI States

Latest PR Test (Base): ❌ Run #31657253640
Latest PR Test (Extra): ❌ Run #31657253406

[gemini-code-assist]

Code Review

This pull request introduces periodic runtime load snapshots (LoadStat) published on a distinct ZMQ topic (LOAD_TOPIC) over the same socket as KV-cache events, allowing load-aware routers to read the true queue depth and KV occupancy. It updates the event publisher to support topic overrides with independent sequence streams and adds a wire-contract test to ensure the msgpack array shape matches expectations. The review feedback suggests using dataclasses.field(init=False, repr=False) for the internal counter fields in SchedulerKvEventsPublisher to avoid exposing them in the autogenerated constructor and representation methods.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Since SchedulerKvEventsPublisher is a dataclass, declaring _load_publish_counter and _load_publish_fail_count as standard fields with default values means they will be included in the autogenerated __init__ signature, __repr__, and __eq__ methods. Since these are internal/private state variables, they should be excluded from the constructor and representation by using dataclasses.field(default=0, init=False, repr=False).

Suggested change

	# Throttle counter for publish_load_stat (slots dataclass: must be declared).
	_load_publish_counter: int = 0
	# Consecutive publish_load_stat failures, reset on success. Drives the
	# periodic re-warn (see LOAD_PUBLISH_FAIL_WARN_EVERY).
	_load_publish_fail_count: int = 0
	# Throttle counter for publish_load_stat (slots dataclass: must be declared).
	_load_publish_counter: int = dataclasses.field(default=0, init=False, repr=False)
	# Consecutive publish_load_stat failures, reset on success. Drives the
	# periodic re-warn (see LOAD_PUBLISH_FAIL_WARN_EVERY).
	_load_publish_fail_count: int = dataclasses.field(default=0, init=False, repr=False)