feat(scheduling): workload- and capacity-aware routing for heterogeneous rollout buckets

[shiinakkk]

Summary

Upgrades heterogeneous rollout scheduling from request-count load balancing to a three-layer signal system, and fixes correctness issues found in an audit of the live rollout call chain.

Problems solved

• Blind load balancing: routing by min(active_requests) ignored request weight and bucket capacity — a 20K-token agentic trajectory counted the same as a 200-token answer, and small TP-1 buckets were compared against TP-4 buckets on raw counts.
• No feedback from vLLM: schedulers never saw real KV occupancy, waiting queues, or preemptions; estimates drifted silently.
• Multi-rank herding: each training rank ranked instances by its own partial view, so N ranks stampeded the same "idle" instance.
• Wasted prefill from scattering: n_samples repeats and multi-turn replays were spread across instances, defeating vLLM's prefix cache.
• Accounting/state bugs: scheduler-failure fallback double-decremented unrelated requests (fooling wait_until_idle), weight-sync rebinding reset learned state every step, kv-ratio selection mixed units, and partial capacity tables silently starved large buckets.

Approach

1. Token-metered accounting — per-request reservation (prompt + EMA-expected output, clamped by max_new_tokens) retired by reservation ID, so out-of-order completions can't steal load from in-flight requests; counters conserve to exactly zero.
2. KV-capacity normalization (load_metric: kv_tokens) — occupancy ratios per instance, with capacity resolved from vLLM's own profiled kv_cache_size_tokens (via /metrics) first, then last profiled value, then analytic/config estimate.
3. Closed-loop /metrics feedback — engine-layer staggered poller (measured <1% latency impact even at 5Hz); additive bias fusion (local + EMA(observed − local)) instead of max(); hysteresis admission (block >0.90, unblock <0.75) with preemption-event TTL (counters are cumulative — events expire, no permanent blacklisting); graceful degradation to open-loop on feed outage.
4. Cross-rank aggregation — each rank publishes per-instance load via atomic-rename files + heartbeat TTL; ranking uses cluster-wide totals (owner-token guarded against rebind races).
5. Prefix affinity — bounded prompt_id → instance LRU so repeated prompts hit warm prefix caches; strictly subordinate to readiness/queue/capacity/admission filters (blocked or full sticky instances are bypassed and remapped). Live validation on a TP1+TP2 cluster: cache hits concentrated +2304/0 on the sticky instance vs +1152/+1152 scattered with affinity off.
6. Exact token counts — all workflows now pass counts they already computed; unknown callers fall back to chars/4.

All features are config-gated and default-off (load_metric: requests preserves legacy behavior). Verified by behavioral regression tests (tests/test_scheduling_audit.py), the scheduler smoke suite, and a reproducible live-HTTP validator (examples/validate_rollout_scheduling.py) with a JSON report of routing decisions and real vLLM cache-counter deltas.