136 lines
4.5 KiB
Rust
136 lines
4.5 KiB
Rust
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
use std::sync::Arc;
|
|
|
|
use sgl_router::buckets_reorg::{Bucket, BucketGroups, BucketResolver, EngineGroup, SloPreference};
|
|
use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy;
|
|
use sgl_router::policies_reorg::PickError;
|
|
use sgl_router::state::load_monitor::engine_reported_load::EngineReportedLoadTable;
|
|
|
|
fn bucket(id: &str, max: u64, rank: u32, ttft: Option<u64>, tps: Option<f64>) -> Bucket {
|
|
let mut bucket = Bucket::new(
|
|
id,
|
|
BucketGroups::Plain(EngineGroup::new(Arc::new(PowerOfTwoPolicy::new(
|
|
EngineReportedLoadTable::new(),
|
|
)))),
|
|
);
|
|
bucket.limits.max = Some(max);
|
|
bucket.rank = rank;
|
|
bucket.ttft_ms = ttft;
|
|
bucket.tokens_per_second = tps;
|
|
bucket
|
|
}
|
|
|
|
fn resolver() -> BucketResolver {
|
|
BucketResolver::new(vec![
|
|
bucket("both", 100, 9, Some(50), Some(100.0)),
|
|
bucket("ttft", 10, 1, Some(50), Some(10.0)),
|
|
bucket("tps", 10, 2, Some(100), Some(100.0)),
|
|
bucket("neither", 8, 0, Some(100), Some(10.0)),
|
|
bucket("missing", 10, 3, None, None),
|
|
bucket("too-short", 1, 0, Some(1), Some(1000.0)),
|
|
])
|
|
.unwrap()
|
|
}
|
|
|
|
fn ids(resolver: &BucketResolver, ttft: Option<u64>, tps: Option<f64>) -> Vec<&str> {
|
|
resolver
|
|
.resolve(5, Some(5), ttft, tps)
|
|
.unwrap()
|
|
.into_iter()
|
|
.map(|b| b.id.as_str())
|
|
.collect()
|
|
}
|
|
|
|
#[test]
|
|
fn slo_tiers_keep_length_constraints_and_capacity_rank_order() {
|
|
let mut resolver = resolver();
|
|
assert_eq!(
|
|
ids(&resolver, Some(50), Some(100.0)),
|
|
["neither", "ttft", "tps", "missing", "both"]
|
|
);
|
|
resolver.ttft_slo = SloPreference::SloFirst;
|
|
resolver.tps_slo = SloPreference::SloFirst;
|
|
assert_eq!(
|
|
ids(&resolver, Some(50), Some(100.0)),
|
|
["both", "ttft", "tps", "neither", "missing"]
|
|
);
|
|
resolver.ttft_slo = SloPreference::BestEffort;
|
|
resolver.tps_slo = SloPreference::BestEffort;
|
|
assert_eq!(
|
|
ids(&resolver, Some(50), Some(100.0)),
|
|
["neither", "missing", "ttft", "tps", "both"]
|
|
);
|
|
resolver.ttft_slo = SloPreference::SloFirst;
|
|
assert_eq!(
|
|
ids(&resolver, Some(50), Some(100.0)),
|
|
["ttft", "neither", "missing", "both", "tps"]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn absent_targets_are_neutral_and_each_preference_can_be_disabled() {
|
|
let mut resolver = resolver();
|
|
resolver.ttft_slo = SloPreference::SloFirst;
|
|
resolver.tps_slo = SloPreference::BestEffort;
|
|
assert_eq!(
|
|
ids(&resolver, None, None),
|
|
["neither", "ttft", "tps", "missing", "both"]
|
|
);
|
|
assert_eq!(
|
|
ids(&resolver, Some(50), None),
|
|
["ttft", "both", "neither", "tps", "missing"]
|
|
);
|
|
resolver.ttft_slo = SloPreference::Disabled;
|
|
resolver.tps_slo = SloPreference::SloFirst;
|
|
assert_eq!(
|
|
ids(&resolver, Some(1), Some(100.0)),
|
|
["tps", "both", "neither", "ttft", "missing"]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn invalid_enabled_targets_fail_and_invalid_estimates_do_not_match() {
|
|
let mut resolver = resolver();
|
|
resolver.ttft_slo = SloPreference::SloFirst;
|
|
resolver.tps_slo = SloPreference::SloFirst;
|
|
assert!(matches!(
|
|
resolver.resolve(5, None, Some(0), None),
|
|
Err(PickError::InvalidSignal(_))
|
|
));
|
|
for tps in [0.0, -1.0, f64::NAN, f64::INFINITY] {
|
|
assert!(matches!(
|
|
resolver.resolve(5, None, None, Some(tps)),
|
|
Err(PickError::InvalidSignal(_))
|
|
));
|
|
}
|
|
resolver.buckets[0].ttft_ms = Some(0);
|
|
resolver.buckets[0].tokens_per_second = Some(f64::INFINITY);
|
|
assert_eq!(
|
|
ids(&resolver, Some(50), Some(100.0)),
|
|
["ttft", "tps", "neither", "missing", "both"]
|
|
);
|
|
resolver.ttft_slo = SloPreference::Disabled;
|
|
resolver.tps_slo = SloPreference::Disabled;
|
|
assert!(resolver.resolve(5, None, Some(0), Some(f64::NAN)).is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn slo_preferences_never_relax_peak_capacity_or_input_range() {
|
|
let mut resolver = resolver();
|
|
resolver.ttft_slo = SloPreference::SloFirst;
|
|
resolver.tps_slo = SloPreference::SloFirst;
|
|
resolver.buckets[0].max_context_tokens = Some(6);
|
|
resolver.buckets[1].limits.min = Some(6);
|
|
let found = resolver.resolve(5, Some(7), Some(50), Some(100.0)).unwrap();
|
|
assert_eq!(
|
|
found.iter().map(|b| b.id.as_str()).collect::<Vec<_>>(),
|
|
["tps", "neither", "missing"]
|
|
);
|
|
assert!(matches!(
|
|
resolver.resolve(5, Some(4), None, None),
|
|
Err(PickError::InvalidSignal(_))
|
|
));
|
|
}
|