[Router] Add composable scoring and eligibility policies (#37731)
Co-authored-by: inkcherry <mingzhi.liu@amd.com>
This commit is contained in:
@@ -133,6 +133,8 @@ async fn static_urls_pd_role_resolved_end_to_end() {
|
||||
circuit_breaker: None,
|
||||
cache_aware: None,
|
||||
sticky: None,
|
||||
fused: None,
|
||||
eligibility: None,
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec![url.clone()],
|
||||
|
||||
@@ -72,6 +72,8 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() {
|
||||
circuit_breaker: None,
|
||||
cache_aware: None,
|
||||
sticky: None,
|
||||
fused: None,
|
||||
eligibility: None,
|
||||
},
|
||||
discovery: sgl_router::config::DiscoveryBackend::StaticUrls(
|
||||
sgl_router::config::StaticUrlsDiscoveryConfig {
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
//! Two REAL scoring policies whose terms DISAGREE, composed and routed. The
|
||||
//! in-crate fusion tests sum `ByIndex` stubs that rank the fleet the SAME way,
|
||||
//! so their `select()` half lands on ws[2] whichever term you read; this one's
|
||||
//! half discriminates. (Their `scores()` half does catch a dropped term —
|
||||
//! verified by mutation, so this file does not claim otherwise.)
|
||||
//!
|
||||
//! NOT pinned here: how `load_based` scales load — W2's min-max scale-free
|
||||
//! defect is unruled, and both candidate curves put the busiest worker at 0.0
|
||||
//! 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::kv_events::{
|
||||
compute_block_hashes, BlockSizeOracle, HashTree, KvWorkerId,
|
||||
};
|
||||
use sgl_router::policies::load_based::LoadBasedPolicy;
|
||||
use sgl_router::policies::scoring::{prefix_cache::PrefixCachePolicy, FusedScorePolicy};
|
||||
use sgl_router::policies::{Policy, SelectionContext};
|
||||
use sgl_router::workers::Worker;
|
||||
use std::sync::Arc;
|
||||
|
||||
const BLOCK: usize = 4;
|
||||
|
||||
fn worker(id: &str) -> Arc<Worker> {
|
||||
Arc::new(Worker::new(WorkerSpec {
|
||||
id: WorkerId(id.into()),
|
||||
url: id.into(),
|
||||
mode: WorkerMode::Plain,
|
||||
model_ids: vec![ModelId("tiny".into())],
|
||||
bootstrap_port: None,
|
||||
}))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_weight_override_steers_a_two_term_fusion_past_either_term_alone() {
|
||||
let ids: Vec<u32> = (0..(BLOCK as u32 * 4)).collect();
|
||||
let tree = Arc::new(HashTree::new());
|
||||
tree.insert(
|
||||
&KvWorkerId::new("hot".into(), 0),
|
||||
None,
|
||||
&compute_block_hashes(&ids, BLOCK),
|
||||
);
|
||||
let oracle = BlockSizeOracle::new();
|
||||
oracle.try_set(BLOCK as u32).expect("fresh oracle");
|
||||
|
||||
// "hot" holds the whole prompt AND is the busiest: the two terms disagree.
|
||||
let ws = vec![worker("hot"), worker("cold")];
|
||||
let _held: Vec<_> = (0..3).map(|_| ws[0].load_guard()).collect();
|
||||
|
||||
let model = ModelId("tiny".into());
|
||||
let ctx = SelectionContext::new(&model, None).with_request_tokens(Some(&ids));
|
||||
let cache = || PrefixCachePolicy::new(Arc::clone(&tree), Arc::clone(&oracle), 1.0);
|
||||
|
||||
// Vacuity guard: if the terms agreed, no weight could change the answer and
|
||||
// everything below would pass against a composer that read only one of them.
|
||||
assert_eq!(cache().select(&ws, &ctx).unwrap().id, ws[0].id, "cache→hot");
|
||||
assert_eq!(
|
||||
LoadBasedPolicy::new().select(&ws, &ctx).unwrap().id,
|
||||
ws[1].id,
|
||||
"load→cold"
|
||||
);
|
||||
|
||||
// Same two terms, same fleet, same request — only the override differs.
|
||||
for (load_weight, want) in [(0.25, &ws[0]), (4.0, &ws[1])] {
|
||||
let fused = FusedScorePolicy::new(vec![
|
||||
(Arc::new(cache()) as Arc<dyn Policy>, None),
|
||||
(Arc::new(LoadBasedPolicy::new()), Some(load_weight)),
|
||||
])
|
||||
.expect("both terms are fusable");
|
||||
let got = fused.select(&ws, &ctx).expect("non-empty fleet");
|
||||
assert_eq!(got.id, want.id, "--fuse load_based={load_weight}");
|
||||
}
|
||||
}
|
||||
@@ -4,6 +4,7 @@
|
||||
mod zmq_helpers;
|
||||
|
||||
mod cache_aware_zmq;
|
||||
mod fused_score;
|
||||
mod kv_events_hash_parity;
|
||||
mod kv_events_tree_concurrent;
|
||||
mod kv_events_two_subscribers;
|
||||
|
||||
@@ -12,10 +12,18 @@
|
||||
//! doesn't render tool schemas, so its ids would diverge from the engine).
|
||||
//! * A request with multimodal (array) content → `input_ids` omitted (a text
|
||||
//! tokenizer can't represent image content).
|
||||
//!
|
||||
//! The model id contains `deepseek-v4` so the tokenizer registry auto-attaches
|
||||
//! the built-in V4 chat encoder — the engine-equivalent path — without a
|
||||
//! template fixture.
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use serde_json::{json, Value};
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, CacheAwareConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig,
|
||||
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use sgl_router::policies::factory::build_registry;
|
||||
use sgl_router::policies::kv_events::{BlockSizeOracle, HashTree};
|
||||
@@ -28,9 +36,35 @@ use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tower::ServiceExt;
|
||||
|
||||
use crate::common::cache_aware_fixture::{config, MODEL};
|
||||
use crate::common::mock_worker::MockWorker;
|
||||
|
||||
const MODEL: &str = "deepseek-v4-tiny";
|
||||
|
||||
fn config() -> Config {
|
||||
Config {
|
||||
server: ServerConfig {
|
||||
host: "0".into(),
|
||||
port: 0,
|
||||
},
|
||||
observability: ObservabilityConfig::default(),
|
||||
model: ModelConfig {
|
||||
id: MODEL.into(),
|
||||
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
|
||||
policy: PolicyKind::CacheAwareZmq,
|
||||
circuit_breaker: None,
|
||||
cache_aware: Some(CacheAwareConfig::default()),
|
||||
sticky: None,
|
||||
fused: None,
|
||||
eligibility: None,
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
}),
|
||||
proxy: ProxyConfig::default(),
|
||||
active_load: ActiveLoadConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn build_ctx(url: String) -> Arc<AppContext> {
|
||||
let cfg = config();
|
||||
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
|
||||
|
||||
@@ -37,6 +37,8 @@ fn config_for(_worker_url: &str) -> Config {
|
||||
circuit_breaker: None,
|
||||
cache_aware: None,
|
||||
sticky: None,
|
||||
fused: None,
|
||||
eligibility: None,
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
|
||||
@@ -29,6 +29,8 @@ pub fn config() -> Config {
|
||||
circuit_breaker: None,
|
||||
cache_aware: Some(CacheAwareConfig::default()),
|
||||
sticky: None,
|
||||
fused: None,
|
||||
eligibility: None,
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
|
||||
@@ -41,6 +41,8 @@ async fn failover_when_one_worker_dies() {
|
||||
}),
|
||||
cache_aware: None,
|
||||
sticky: None,
|
||||
fused: None,
|
||||
eligibility: None,
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec![w1.url.clone(), w2.url.clone(), w3.url.clone()],
|
||||
|
||||
@@ -47,6 +47,8 @@ fn build_ctx_with_worker(worker_url: &str) -> Arc<AppContext> {
|
||||
circuit_breaker: None,
|
||||
cache_aware: None,
|
||||
sticky: None,
|
||||
fused: None,
|
||||
eligibility: None,
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
|
||||
@@ -34,6 +34,8 @@ async fn forwards_whitelisted_headers_strips_others() {
|
||||
circuit_breaker: None,
|
||||
cache_aware: None,
|
||||
sticky: None,
|
||||
fused: None,
|
||||
eligibility: None,
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
|
||||
@@ -49,6 +49,8 @@ fn config() -> Config {
|
||||
circuit_breaker: None,
|
||||
cache_aware: None,
|
||||
sticky: None,
|
||||
fused: None,
|
||||
eligibility: None,
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
|
||||
@@ -48,6 +48,8 @@ fn config() -> Config {
|
||||
circuit_breaker: None,
|
||||
cache_aware: None,
|
||||
sticky: None,
|
||||
fused: None,
|
||||
eligibility: None,
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
|
||||
@@ -45,6 +45,8 @@ fn config() -> Config {
|
||||
circuit_breaker: None,
|
||||
cache_aware: None,
|
||||
sticky: None,
|
||||
fused: None,
|
||||
eligibility: None,
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
|
||||
@@ -25,7 +25,7 @@ use axum::http::{Request, StatusCode};
|
||||
use serde_json::{json, Value};
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry;
|
||||
@@ -60,10 +60,12 @@ fn config() -> Config {
|
||||
// mid-test; round-robin fallback for the initial pin of a key.
|
||||
sticky: Some(StickyConfig {
|
||||
header_name: HEADER.to_string(),
|
||||
fallback_policy: PolicyKind::RoundRobin,
|
||||
fallback_policy: StickyFallbackKind::RoundRobin,
|
||||
idle_secs: 3600,
|
||||
eviction_interval_secs: 3600,
|
||||
}),
|
||||
fused: None,
|
||||
eligibility: None,
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
|
||||
use sgl_router::config::{
|
||||
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig,
|
||||
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind,
|
||||
};
|
||||
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry;
|
||||
@@ -46,10 +46,12 @@ fn build_sticky_ctx(header_name: &str, worker_urls: &[String]) -> Arc<AppContext
|
||||
cache_aware: None,
|
||||
sticky: Some(StickyConfig {
|
||||
header_name: header_name.to_string(),
|
||||
fallback_policy: PolicyKind::RoundRobin,
|
||||
fallback_policy: StickyFallbackKind::RoundRobin,
|
||||
idle_secs: 3600,
|
||||
eviction_interval_secs: 3600,
|
||||
}),
|
||||
fused: None,
|
||||
eligibility: None,
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
|
||||
@@ -41,6 +41,8 @@ fn config(_worker_url: &str) -> Config {
|
||||
circuit_breaker: None,
|
||||
cache_aware: None,
|
||||
sticky: None,
|
||||
fused: None,
|
||||
eligibility: None,
|
||||
},
|
||||
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
|
||||
urls: vec!["http://placeholder:0".into()],
|
||||
|
||||
Reference in New Issue
Block a user