[sgl-router] refactor - move policy-required states under src/state (#40272)

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Kan Wu
2026-09-20 16:49:23 -07:00
committed by GitHub
co-authored by Claude Fable 5.1
parent acd20a516e
commit 4a9dc5c4af
78 changed files with 787 additions and 655 deletions
@@ -5,10 +5,8 @@ use std::collections::HashMap;
use std::sync::Arc;
use sgl_kv_indexer::PrefixOutcome;
use sgl_router::policies::kv_events::{
compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId,
};
use sgl_router::policies::prefix_provider::RadixTreePrefixProvider;
use sgl_router::state::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId};
#[test]
fn radix_tree_reports_contiguous_prefix_depth_per_worker() {
@@ -12,8 +12,10 @@ use sgl_router::policies::decode::{
resolve_decode_with_capacity_fallback, DecodePolicy, DecodePowerOfTwoPolicy,
DecodeSelectionContext, LegacyHostAffinityDecodePolicy,
};
use sgl_router::policies::engine_load::{EngineLoadSnapshot, NativeCacheWorkerLoad};
use sgl_router::policies::SelectionProposal;
use sgl_router::state::load_monitor::engine_reported_load::{
EngineReportedLoadSnapshot, EngineReportedSchedulingLoad,
};
use sgl_router::workers::Worker;
use std::collections::HashMap;
use std::sync::atomic::Ordering;
@@ -30,15 +32,15 @@ fn worker(id: &str) -> Arc<Worker> {
}))
}
fn snapshot(entries: &[(&Arc<Worker>, u64, u64, u64, u64)]) -> EngineLoadSnapshot {
EngineLoadSnapshot::from_native_cache_workers(
fn snapshot(entries: &[(&Arc<Worker>, u64, u64, u64, u64)]) -> EngineReportedLoadSnapshot {
EngineReportedLoadSnapshot::from_native_cache_workers(
7,
entries
.iter()
.map(|(worker, running, waiting, used, capacity)| {
(
worker.url.clone(),
NativeCacheWorkerLoad {
EngineReportedSchedulingLoad {
num_running_reqs: *running,
num_waiting_reqs: *waiting,
num_waiting_uncached_tokens: *waiting,
@@ -12,15 +12,15 @@
//! and the idlest at 1.0, so every assertion below holds either way.
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::engine_load::{EngineLoadSnapshot, EngineWorkerLoad};
use sgl_router::policies::kv_events::{
compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId,
};
use sgl_router::policies::load_based::LoadBasedPolicy;
use sgl_router::policies::scoring::{
prefix_cache::PrefixCachePolicy, FusedScorePolicy, ScorePolicy,
};
use sgl_router::policies::{Policy, SelectionContext};
use sgl_router::state::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId};
use sgl_router::state::load_monitor::engine_reported_load::{
EngineReportedLoadSnapshot, EngineReportedWorkerLoad,
};
use sgl_router::workers::Worker;
use std::{collections::HashMap, sync::Arc, time::Instant};
@@ -82,12 +82,12 @@ fn fused_load_based_term_uses_the_request_snapshot() {
let ws = vec![worker("w0"), worker("w1")];
// Local counters changed after the request snapshot and prefer w0.
let _after_snapshot: Vec<_> = (0..10).map(|_| ws[1].load_guard()).collect();
let snapshot = EngineLoadSnapshot::from_workers(
let snapshot = EngineReportedLoadSnapshot::from_workers(
29,
HashMap::from([
(
ws[0].url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 50,
num_waiting_reqs: 0,
num_tokens: 0,
@@ -97,7 +97,7 @@ fn fused_load_based_term_uses_the_request_snapshot() {
),
(
ws[1].url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 1,
num_waiting_reqs: 0,
num_tokens: 0,
@@ -123,12 +123,12 @@ fn fused_load_based_term_uses_the_request_snapshot() {
fn score_policy_forwards_the_request_snapshot_to_load_based() {
let ws = vec![worker("w0"), worker("w1")];
let _after_snapshot: Vec<_> = (0..10).map(|_| ws[1].load_guard()).collect();
let snapshot = EngineLoadSnapshot::from_workers(
let snapshot = EngineReportedLoadSnapshot::from_workers(
31,
HashMap::from([
(
ws[0].url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 50,
num_waiting_reqs: 0,
num_tokens: 0,
@@ -138,7 +138,7 @@ fn score_policy_forwards_the_request_snapshot_to_load_based() {
),
(
ws[1].url.clone(),
EngineWorkerLoad {
EngineReportedWorkerLoad {
num_running_reqs: 1,
num_waiting_reqs: 0,
num_tokens: 0,
@@ -19,7 +19,7 @@
//! whatever fixture is checked in.
use serde::Deserialize;
use sgl_router::policies::kv_events::compute_block_hashes;
use sgl_router::state::kv_events::compute_block_hashes;
use std::path::PathBuf;
#[derive(Debug, Deserialize)]
@@ -3,7 +3,7 @@
//! Concurrent-mutation stress test for `HashTree`.
//!
//! The inline tests in `policies::kv_events::tree` are all
//! The inline tests in `state::kv_events::tree` are all
//! single-threaded. Under production load, multiple worker subscribers
//! drive `insert` / `remove` / `clear_worker` against the same tree from
//! tokio worker threads while the chat handler simultaneously calls
@@ -25,7 +25,7 @@
use std::sync::Arc;
use std::thread;
use sgl_router::policies::kv_events::{HashTree, KvWorkerId};
use sgl_router::state::kv_events::{HashTree, KvWorkerId};
fn worker(i: usize) -> KvWorkerId {
KvWorkerId {
@@ -22,8 +22,8 @@ use std::time::Duration;
use zeromq::SocketSend;
use sgl_router::policies::kv_events::discovery::EventConfig;
use sgl_router::policies::kv_events::{compute_block_hashes, KvEventIndex, KvWorkerId};
use sgl_router::state::kv_events::discovery::EventConfig;
use sgl_router::state::kv_events::{compute_block_hashes, KvEventIndex, KvWorkerId};
use super::zmq_helpers::{
build_multipart, encode_block_stored_event, encode_event_batch, make_pub_bound,
@@ -1,7 +1,7 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Shared ZMQ wire-format helpers for the `policies::kv_events` component
//! Shared ZMQ wire-format helpers for the `state::kv_events` component
//! tests. Encodes events in the same msgspec layout SGLang emits, builds
//! the two-frame `[seq, payload]` ZMQ message a real publisher sends, and
//! binds a loopback PUB socket on an OS-assigned port.