feat: add push-based engine load reporting and Router load monitor

[Bakerjc-bgner]

Motivation

Load-aware routing needs fresh, engine-authoritative data. Polling coarse worker-level endpoints cannot reliably expose short-lived queue changes or load imbalance hidden across data-parallel ranks, and independently refreshed values can give one routing decision a torn view of the cluster.

This PR introduces an end-to-end, push-based load-control plane between SGLang engines and experimental/sgl-router. Engines continuously stream per-DP-rank load over long-lived gRPC connections, while the Router validates, aggregates, and publishes immutable snapshots for diagnostics and load-aware policies.

The design provides fine-grained visibility without putting polling fan-out on the request path. It is efficient, topology-aware, resilient to stale or out-of-order data, and opt-in for backward compatibility.

Modifications

Architecture

flowchart LR
    subgraph Engine["SGLang Engine"]
        Scheduler["Scheduler<br/>per-DP-rank load metrics"]
        Snapshot["SHM / ZMQ load snapshot"]
        Reporter["LoadReporterRuntime<br/>single owner"]
        Workers["Tokenizer workers<br/>coalesced IPC control + refresh"]

        Scheduler --> Snapshot --> Reporter
        Workers -. "multi-tokenizer mode" .-> Reporter
    end

    subgraph Router["experimental/sgl-router"]
        Discovery["Worker discovery"]
        Registration["Registration + lease renewal"]
        Ingest["LoadMonitorService.Report<br/>client-streaming gRPC"]
        Validation["Identity, sequence,<br/>and freshness validation"]
        Aggregate["Per-DP-rank aggregation"]
        View["Immutable, versioned snapshot"]
        Diagnostic["GET /v1/load_monitor/snapshot"]
        Policy["PolicyCandidate<br/>load-aware routing"]

        Discovery --> Registration
        Ingest --> Validation --> Aggregate --> View
        View --> Diagnostic
        View --> Policy
    end

    Registration -->|"POST /v1/start_reporting<br/>target + interval + lease"| Reporter
    Reporter -->|"long-lived stream<br/>ms-resolution report cadence"| Ingest

Loading

Engine-side load reporter

• Add scheduler-owned load snapshots containing authoritative per-DP-rank queue, token, capacity, cache, utilization, generation-throughput, and prefill-throughput metrics.
• Add a composable reporter runtime with sampling, latest-snapshot storage, report construction, fixed-rate scheduling, immediate first delivery, retry/backoff, lease renewal, and bounded shutdown.
• Use one persistent client-streaming gRPC connection per Router target. The reporting interval is expressed in milliseconds, allowing the control plane to capture fast load changes without repeatedly creating connections.
• Preserve single ownership in both tokenizer topologies: the HTTP process owns the reporter in single-tokenizer mode; multi-tokenizer workers use IPC control and coalesced refresh notifications instead of opening duplicate streams.
• Add the reporter as an optional Python dependency and keep generated protobuf files protected from incompatible formatter rewrites.

Router-side load monitor

• Add an opt-in gRPC load-monitor listener and worker registration loop. The Router tells each discovered engine where and how often to report through POST /v1/start_reporting and renews the reporting lease.
• Validate immutable stream identity, monotonically increasing sequence IDs, rank payloads, and Router-observed freshness before accepting data. Stale and out-of-order reports cannot silently influence routing.
• Aggregate per-DP-rank reports into worker-level capacity, free-token, available-slot, queue-pressure, request-utilization, weighted token-usage, maximum-rank token-usage, and throughput signals. This preserves hotspot information that a coarse worker average would hide.
• Publish fully owned, versioned snapshots under one store lock. Each request captures exactly one snapshot for both prefill and decode selection, preventing mixed-version routing decisions while reports arrive concurrently.
• Expose a read-only GET /v1/load_monitor/snapshot diagnostic endpoint and attach fresh aggregate load to policy candidates, including the existing load-based policy path.
• Add a shared-protobuf parity test so the Python engine and Rust Router cannot drift unnoticed.

Why this design

• Fresh and precise: engine-pushed, per-rank signals capture queue and capacity changes at millisecond-resolution cadences and avoid masking a hot DP rank behind a worker average.
• Efficient: long-lived streams reuse connections and remove Router-side polling fan-out from the hot path.
• Consistent: immutable, per-request snapshots give routing policies one coherent cluster view.
• Resilient: identity, sequence, freshness, retry, and lease semantics prevent stale workers, delayed reports, or process reincarnations from corrupting decisions.
• Extensible: the versioned snapshot is a clean, read-only foundation for additional load-balancing policies and observability without coupling policy code to transport details.
• Compatible: monitoring is opt-in, reporter dependencies are optional, and the existing routing behavior remains unchanged when load monitoring is disabled.

Accuracy Tests

N/A for model accuracy. This PR does not change model forward code, kernels, sampling, or generated outputs.

Functional and correctness validation:

• Built a container image and completed deployment smoke testing for Engine registration, continuous load reporting, Router ingestion, and fresh snapshot output.
• cargo test --lib load_monitor: 19 passed, 0 failed.
• cargo test --lib policies::: 175 passed, 0 failed.
• Added Python unit and registered-test coverage for snapshot backends, reporter lifecycle, multi-tokenizer single ownership, IPC round trips, prefill-throughput accounting, and multi-worker reporting.
• Added Rust coverage for report validation, identity and sequence handling, freshness, aggregation, immutable snapshots, policy integration, worker registration/retry/removal, fake-engine gRPC reporting, and Python/Rust protobuf parity.

Speed Tests and Profiling

N/A for inference throughput benchmarking. The feature is an opt-in control-plane path and does not modify model execution or the inference hot path.

The transport is intentionally lightweight: it reuses long-lived gRPC streams, samples already-published engine snapshots, and coalesces multi-tokenizer refresh notifications. No inference-speed improvement is claimed by this PR.

Additional local checks:

• cargo fmt --all -- --check: passed.
• cargo check --all-targets: passed.
• Python 3.12 syntax compilation for all 30 changed Python files: passed.
• git diff --check: passed.
• Full GitHub CI: to be triggered after the PR is opened.

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.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ⏳ Run #30255759319
Latest PR Test (Extra): ⏳ Run #30255759497

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[Bakerjc-bgner]

Hi @hzh0425 and @whybeyoung , would you mind taking a look at this PR? Thank you!

[Bakerjc-bgner]

Superseded by #32523, which contains the Engine Load Reporter only and is split into reviewable incremental commits. The Router Load Monitor and Prefill throughput extension will be submitted separately.