From ae68158f2e7c56af09d8dabca2eeaf40804ecdb9 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Wed, 12 Nov 2025 16:36:04 +0900 Subject: [PATCH] [router] move radix tree to policy crate and addreses some code styles (#13131) --- sgl-router/src/lib.rs | 1 - sgl-router/src/policies/bucket.rs | 14 +++---- sgl-router/src/policies/cache_aware.rs | 52 +++++++++++++------------- sgl-router/src/policies/mod.rs | 1 + sgl-router/src/{ => policies}/tree.rs | 6 ++- 5 files changed, 39 insertions(+), 35 deletions(-) rename sgl-router/src/{ => policies}/tree.rs (99%) diff --git a/sgl-router/src/lib.rs b/sgl-router/src/lib.rs index 2af3561e3..8dfcc083e 100644 --- a/sgl-router/src/lib.rs +++ b/sgl-router/src/lib.rs @@ -19,7 +19,6 @@ pub mod server; pub mod service_discovery; pub mod tokenizer; pub mod tool_parser; -pub mod tree; use crate::metrics::PrometheusConfig; #[pyclass(eq)] diff --git a/sgl-router/src/policies/bucket.rs b/sgl-router/src/policies/bucket.rs index c491c5c95..ea4a3a440 100644 --- a/sgl-router/src/policies/bucket.rs +++ b/sgl-router/src/policies/bucket.rs @@ -672,7 +672,7 @@ mod tests { .select_worker(&prefill_workers, Some(&*"a".repeat(34))) .unwrap(); - tokio::time::sleep(Duration::from_secs(10)).await; + tokio::time::sleep(Duration::from_secs(11)).await; { let model_key = "default"; @@ -844,7 +844,7 @@ mod tests { .select_worker(&prefill_workers, Some(&*"a".repeat(26))) .unwrap(); - tokio::time::sleep(Duration::from_secs(3)).await; + tokio::time::sleep(Duration::from_secs(4)).await; // Verify boundaries adjusted to: [0, 20], [21, 26], [27, MAX] { let model_key = "default"; @@ -889,7 +889,7 @@ mod tests { .select_worker(&prefill_workers, Some(&*"a".repeat(57))) .unwrap(); - tokio::time::sleep(Duration::from_secs(3)).await; + tokio::time::sleep(Duration::from_secs(4)).await; // Verify boundaries adjusted to: [0, 40], [41, 57], [58, MAX] { let model_key = "default"; @@ -974,7 +974,7 @@ mod tests { .select_worker(&prefill_workers, Some(&*"a".repeat(20))) .unwrap(); - tokio::time::sleep(Duration::from_secs(3)).await; + tokio::time::sleep(Duration::from_secs(4)).await; { let model_key = "default"; @@ -1001,7 +1001,7 @@ mod tests { .select_worker(&prefill_workers, Some(&*"a".repeat(7))) .unwrap(); - tokio::time::sleep(Duration::from_secs(3)).await; + tokio::time::sleep(Duration::from_secs(4)).await; { let model_key = "default"; @@ -1099,7 +1099,7 @@ mod tests { .select_worker(&prefill_workers, Some(&*"a".repeat(26))) .unwrap(); - tokio::time::sleep(Duration::from_secs(3)).await; + tokio::time::sleep(Duration::from_secs(4)).await; { let model_key = "default"; @@ -1141,7 +1141,7 @@ mod tests { .select_worker(&prefill_workers, Some(&*"a".repeat(55))) .unwrap(); - tokio::time::sleep(Duration::from_secs(3)).await; + tokio::time::sleep(Duration::from_secs(4)).await; { let model_key = "default"; diff --git a/sgl-router/src/policies/cache_aware.rs b/sgl-router/src/policies/cache_aware.rs index 99fbc1558..c43d34ed8 100644 --- a/sgl-router/src/policies/cache_aware.rs +++ b/sgl-router/src/policies/cache_aware.rs @@ -65,8 +65,8 @@ use dashmap::DashMap; use rand::Rng; use tracing::debug; -use super::{get_healthy_worker_indices, CacheAwareConfig, LoadBalancingPolicy}; -use crate::{core::Worker, metrics::RouterMetrics, tree::Tree}; +use super::{get_healthy_worker_indices, tree::Tree, CacheAwareConfig, LoadBalancingPolicy}; +use crate::{core::Worker, metrics::RouterMetrics}; /// Cache-aware routing policy /// @@ -346,30 +346,6 @@ impl LoadBalancingPolicy for CacheAwarePolicy { } } - fn name(&self) -> &'static str { - "cache_aware" - } - - fn needs_request_text(&self) -> bool { - true // Cache-aware policy needs request text for cache affinity - } - - fn on_request_complete(&self, worker_url: &str, success: bool) { - // Could track success rates per worker for more intelligent routing - if !success { - // Optionally reduce affinity for failed requests - tracing::debug!( - "Request to {} completed with success={}", - worker_url, - success - ); - } - } - - fn as_any(&self) -> &dyn std::any::Any { - self - } - fn select_worker_pair( &self, prefill_workers: &[Arc], @@ -400,6 +376,30 @@ impl LoadBalancingPolicy for CacheAwarePolicy { Some((prefill_idx, decode_idx)) } + + fn on_request_complete(&self, worker_url: &str, success: bool) { + // Could track success rates per worker for more intelligent routing + if !success { + // Optionally reduce affinity for failed requests + tracing::debug!( + "Request to {} completed with success={}", + worker_url, + success + ); + } + } + + fn name(&self) -> &'static str { + "cache_aware" + } + + fn needs_request_text(&self) -> bool { + true // Cache-aware policy needs request text for cache affinity + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } } impl Default for CacheAwarePolicy { diff --git a/sgl-router/src/policies/mod.rs b/sgl-router/src/policies/mod.rs index 67764eebb..7eca86097 100644 --- a/sgl-router/src/policies/mod.rs +++ b/sgl-router/src/policies/mod.rs @@ -14,6 +14,7 @@ mod power_of_two; mod random; mod registry; mod round_robin; +mod tree; pub use bucket::BucketPolicy; pub use cache_aware::CacheAwarePolicy; diff --git a/sgl-router/src/tree.rs b/sgl-router/src/policies/tree.rs similarity index 99% rename from sgl-router/src/tree.rs rename to sgl-router/src/policies/tree.rs index e4661620c..4ad10fff1 100644 --- a/sgl-router/src/tree.rs +++ b/sgl-router/src/policies/tree.rs @@ -321,7 +321,7 @@ impl Tree { (ret_text, tenant) } - #[allow(unused_assignments)] + #[allow(unused_assignments, dead_code)] pub fn prefix_match_tenant(&self, text: &str, tenant: &str) -> String { let mut curr = Arc::clone(&self.root); let mut curr_idx = 0; @@ -529,6 +529,7 @@ impl Tree { self.tenant_char_count.remove(&tenant.to_string()); } + #[allow(dead_code)] pub fn get_tenant_char_count(&self) -> HashMap { self.tenant_char_count .iter() @@ -560,6 +561,7 @@ impl Tree { min_tenant.unwrap_or_else(|| "empty".to_string()) } + #[allow(dead_code)] pub fn get_used_size_per_tenant(&self) -> HashMap { // perform a DFS to traverse all nodes and calculate the total size used by each tenant @@ -584,6 +586,7 @@ impl Tree { used_size_per_tenant } + #[allow(dead_code)] fn node_to_string(node: &NodeRef, prefix: &str, is_last: bool) -> String { let mut result = String::new(); @@ -641,6 +644,7 @@ impl Tree { result } + #[allow(dead_code)] pub fn pretty_print(&self) { if self.root.children.is_empty() { return;