[sgl-router] refactor - generalized admission policy definitions (#40271)

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Kan Wu
2026-09-21 16:26:30 +08:00
committed by GitHub
co-authored by Claude Fable 5.1
parent 70b5b03e78
commit a9871012ac
12 changed files with 445 additions and 172 deletions
@@ -10,7 +10,7 @@ use sgl_router::buckets_reorg::{Bucket, BucketGroups, BucketResolver, EngineGrou
use sgl_router::config::AffinityConfig;
use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec};
use sgl_router::policies::prefix_provider::RadixTreePrefixProvider;
use sgl_router::policies_reorg::admission::{Decision, EngineAdmission};
use sgl_router::policies_reorg::admission::{Decision, EngineAdmission, EngineMetrics};
use sgl_router::policies_reorg::cache_aware::{CacheAwarePolicy, CacheSource, PrefixMemo};
use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy;
use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage};
@@ -18,7 +18,7 @@ use sgl_router::state::kv_events::{
compute_block_hashes, compute_block_hashes_bigram, BlockSizeOracle, HashTree, KvWorkerId,
};
use sgl_router::state::load_monitor::engine_reported_load::{
EngineReportedLoadTable, EngineReportedWorkerLoad, LoadStat, NativeCacheRankLoad,
EngineReportedLoadTable, LoadStat, NativeCacheRankLoad,
};
use sgl_router::workers::Worker;
@@ -116,16 +116,11 @@ impl Reject {
}
impl EngineAdmission for Reject {
fn check(
&self,
engine: &Worker,
_: &PickRequest<'_>,
load: Option<&EngineReportedWorkerLoad>,
) -> Result<Decision, PickError> {
fn check(&self, engine: &Worker, metrics: &EngineMetrics) -> Result<Decision, PickError> {
self.calls
.lock()
.unwrap()
.push((engine.id.0.clone(), load.map(|load| load.num_waiting_reqs)));
.push((engine.id.0.clone(), metrics.waiting_requests));
Ok(if engine.id.0 == self.id {
Decision::Reject("full".into())
} else {