[router] allow one router to support different model families and serving mode (#10244)
This commit is contained in:
Simo Lin
333
sgl-router/src/policies/registry.rs
Normal file
333
sgl-router/src/policies/registry.rs
Normal file
@@ -0,0 +1,333 @@
|
||||
/// Policy Registry for managing model-to-policy mappings
|
||||
///
|
||||
/// This registry manages the dynamic assignment of load balancing policies to models.
|
||||
/// When the first worker of a new model is added, it determines the policy for that model.
|
||||
/// All subsequent workers of the same model use the established policy.
|
||||
/// When the last worker of a model is removed, the policy mapping is cleaned up.
|
||||
use super::{
|
||||
CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, PowerOfTwoPolicy, RandomPolicy,
|
||||
RoundRobinPolicy,
|
||||
};
|
||||
use crate::config::types::PolicyConfig;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
/// Registry for managing model-to-policy mappings
|
||||
#[derive(Clone)]
|
||||
pub struct PolicyRegistry {
|
||||
/// Model ID -> Policy instance mapping
|
||||
model_policies: Arc<RwLock<HashMap<String, Arc<dyn LoadBalancingPolicy>>>>,
|
||||
|
||||
/// Model ID -> Worker count for cleanup tracking
|
||||
model_worker_counts: Arc<RwLock<HashMap<String, usize>>>,
|
||||
|
||||
/// Default policy instance (cached)
|
||||
default_policy: Arc<dyn LoadBalancingPolicy>,
|
||||
|
||||
/// Prefill policy for PD mode
|
||||
prefill_policy: Arc<RwLock<Option<Arc<dyn LoadBalancingPolicy>>>>,
|
||||
|
||||
/// Decode policy for PD mode
|
||||
decode_policy: Arc<RwLock<Option<Arc<dyn LoadBalancingPolicy>>>>,
|
||||
}
|
||||
|
||||
impl PolicyRegistry {
|
||||
/// Create a new PolicyRegistry with a default policy
|
||||
pub fn new(default_policy_config: PolicyConfig) -> Self {
|
||||
let default_policy = Self::create_policy_from_config(&default_policy_config);
|
||||
|
||||
Self {
|
||||
model_policies: Arc::new(RwLock::new(HashMap::new())),
|
||||
model_worker_counts: Arc::new(RwLock::new(HashMap::new())),
|
||||
default_policy,
|
||||
prefill_policy: Arc::new(RwLock::new(None)),
|
||||
decode_policy: Arc::new(RwLock::new(None)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Called when a worker is added
|
||||
/// Returns the policy that should be used for this worker's model
|
||||
pub fn on_worker_added(
|
||||
&self,
|
||||
model_id: &str,
|
||||
policy_hint: Option<&str>,
|
||||
) -> Arc<dyn LoadBalancingPolicy> {
|
||||
// Increment worker count
|
||||
{
|
||||
let mut counts = self.model_worker_counts.write().unwrap();
|
||||
*counts.entry(model_id.to_string()).or_insert(0) += 1;
|
||||
debug!(
|
||||
"Worker added for model {}, count: {}",
|
||||
model_id,
|
||||
counts.get(model_id).unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
// Check if model already has a policy
|
||||
{
|
||||
let policies = self.model_policies.read().unwrap();
|
||||
if let Some(existing_policy) = policies.get(model_id) {
|
||||
debug!(
|
||||
"Model {} already has policy: {}",
|
||||
model_id,
|
||||
existing_policy.name()
|
||||
);
|
||||
return Arc::clone(existing_policy);
|
||||
}
|
||||
}
|
||||
|
||||
// New model - determine policy
|
||||
let policy = self.determine_policy_for_model(model_id, policy_hint);
|
||||
|
||||
info!(
|
||||
"Assigning policy {} to new model {}",
|
||||
policy.name(),
|
||||
model_id
|
||||
);
|
||||
|
||||
// Store policy for this model
|
||||
{
|
||||
let mut policies = self.model_policies.write().unwrap();
|
||||
policies.insert(model_id.to_string(), Arc::clone(&policy));
|
||||
}
|
||||
|
||||
policy
|
||||
}
|
||||
|
||||
/// Called when a worker is removed
|
||||
pub fn on_worker_removed(&self, model_id: &str) {
|
||||
let should_cleanup = {
|
||||
let mut counts = self.model_worker_counts.write().unwrap();
|
||||
if let Some(count) = counts.get_mut(model_id) {
|
||||
*count = count.saturating_sub(1);
|
||||
debug!("Worker removed for model {}, count: {}", model_id, *count);
|
||||
if *count == 0 {
|
||||
counts.remove(model_id);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
} else {
|
||||
warn!(
|
||||
"Attempted to remove worker for model {} with no registered workers",
|
||||
model_id
|
||||
);
|
||||
false
|
||||
}
|
||||
};
|
||||
|
||||
// Clean up policy if this was the last worker
|
||||
if should_cleanup {
|
||||
let mut policies = self.model_policies.write().unwrap();
|
||||
if let Some(policy) = policies.remove(model_id) {
|
||||
info!(
|
||||
"Removed policy {} for model {} (last worker removed)",
|
||||
policy.name(),
|
||||
model_id
|
||||
);
|
||||
// Policy will be dropped here, cleaning up any resources
|
||||
drop(policy);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the policy for a model
|
||||
pub fn get_policy(&self, model_id: &str) -> Option<Arc<dyn LoadBalancingPolicy>> {
|
||||
self.model_policies.read().unwrap().get(model_id).cloned()
|
||||
}
|
||||
|
||||
/// Get the default policy
|
||||
pub fn get_default_policy(&self) -> Arc<dyn LoadBalancingPolicy> {
|
||||
Arc::clone(&self.default_policy)
|
||||
}
|
||||
|
||||
/// Get policy for a model, or default if not found
|
||||
pub fn get_policy_or_default(&self, model_id: &str) -> Arc<dyn LoadBalancingPolicy> {
|
||||
self.get_policy(model_id)
|
||||
.unwrap_or_else(|| self.get_default_policy())
|
||||
}
|
||||
|
||||
/// Determine policy for a new model
|
||||
fn determine_policy_for_model(
|
||||
&self,
|
||||
model_id: &str,
|
||||
policy_hint: Option<&str>,
|
||||
) -> Arc<dyn LoadBalancingPolicy> {
|
||||
// 1. Check policy hint from worker
|
||||
if let Some(policy_type) = policy_hint {
|
||||
debug!("Using policy hint '{}' for model {}", policy_type, model_id);
|
||||
return self.create_policy_from_type(policy_type);
|
||||
}
|
||||
|
||||
// 2. Use default policy
|
||||
debug!("Using default policy for model {}", model_id);
|
||||
Arc::clone(&self.default_policy)
|
||||
}
|
||||
|
||||
/// Create a policy from a type string
|
||||
fn create_policy_from_type(&self, policy_type: &str) -> Arc<dyn LoadBalancingPolicy> {
|
||||
match policy_type {
|
||||
"round_robin" => Arc::new(RoundRobinPolicy::new()),
|
||||
"random" => Arc::new(RandomPolicy::new()),
|
||||
"cache_aware" => Arc::new(CacheAwarePolicy::new()),
|
||||
"power_of_two" => Arc::new(PowerOfTwoPolicy::new()),
|
||||
_ => {
|
||||
warn!("Unknown policy type '{}', using default", policy_type);
|
||||
Arc::clone(&self.default_policy)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a policy from a PolicyConfig
|
||||
fn create_policy_from_config(config: &PolicyConfig) -> Arc<dyn LoadBalancingPolicy> {
|
||||
match config {
|
||||
PolicyConfig::RoundRobin => Arc::new(RoundRobinPolicy::new()),
|
||||
PolicyConfig::Random => Arc::new(RandomPolicy::new()),
|
||||
PolicyConfig::CacheAware {
|
||||
cache_threshold,
|
||||
balance_abs_threshold,
|
||||
balance_rel_threshold,
|
||||
eviction_interval_secs,
|
||||
max_tree_size,
|
||||
} => {
|
||||
let cache_config = CacheAwareConfig {
|
||||
cache_threshold: *cache_threshold,
|
||||
balance_abs_threshold: *balance_abs_threshold,
|
||||
balance_rel_threshold: *balance_rel_threshold,
|
||||
eviction_interval_secs: *eviction_interval_secs,
|
||||
max_tree_size: *max_tree_size,
|
||||
};
|
||||
Arc::new(CacheAwarePolicy::with_config(cache_config))
|
||||
}
|
||||
PolicyConfig::PowerOfTwo { .. } => Arc::new(PowerOfTwoPolicy::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get current model->policy mappings (for debugging/monitoring)
|
||||
pub fn get_all_mappings(&self) -> HashMap<String, String> {
|
||||
let policies = self.model_policies.read().unwrap();
|
||||
policies
|
||||
.iter()
|
||||
.map(|(model, policy)| (model.clone(), policy.name().to_string()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Get worker counts per model
|
||||
pub fn get_worker_counts(&self) -> HashMap<String, usize> {
|
||||
self.model_worker_counts.read().unwrap().clone()
|
||||
}
|
||||
|
||||
/// Clear all policies (useful for testing)
|
||||
pub fn clear(&self) {
|
||||
let mut policies = self.model_policies.write().unwrap();
|
||||
policies.clear();
|
||||
let mut counts = self.model_worker_counts.write().unwrap();
|
||||
counts.clear();
|
||||
}
|
||||
|
||||
/// Set the prefill policy for PD mode
|
||||
pub fn set_prefill_policy(&self, policy: Arc<dyn LoadBalancingPolicy>) {
|
||||
let mut prefill_policy = self.prefill_policy.write().unwrap();
|
||||
*prefill_policy = Some(policy);
|
||||
}
|
||||
|
||||
/// Set the decode policy for PD mode
|
||||
pub fn set_decode_policy(&self, policy: Arc<dyn LoadBalancingPolicy>) {
|
||||
let mut decode_policy = self.decode_policy.write().unwrap();
|
||||
*decode_policy = Some(policy);
|
||||
}
|
||||
|
||||
/// Get the prefill policy for PD mode, or default if not set
|
||||
pub fn get_prefill_policy(&self) -> Arc<dyn LoadBalancingPolicy> {
|
||||
let prefill_policy = self.prefill_policy.read().unwrap();
|
||||
prefill_policy
|
||||
.as_ref()
|
||||
.map(Arc::clone)
|
||||
.unwrap_or_else(|| self.get_default_policy())
|
||||
}
|
||||
|
||||
/// Get the decode policy for PD mode, or default if not set
|
||||
pub fn get_decode_policy(&self) -> Arc<dyn LoadBalancingPolicy> {
|
||||
let decode_policy = self.decode_policy.read().unwrap();
|
||||
decode_policy
|
||||
.as_ref()
|
||||
.map(Arc::clone)
|
||||
.unwrap_or_else(|| self.get_default_policy())
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for PolicyRegistry {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("PolicyRegistry")
|
||||
.field("model_policies", &self.model_policies)
|
||||
.field("model_worker_counts", &self.model_worker_counts)
|
||||
.field("default_policy", &self.default_policy.name())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_policy_registry_basic() {
|
||||
let registry = PolicyRegistry::new(PolicyConfig::RoundRobin);
|
||||
|
||||
// First worker of a model sets the policy
|
||||
let policy1 = registry.on_worker_added("llama-3", Some("cache_aware"));
|
||||
assert_eq!(policy1.name(), "cache_aware");
|
||||
|
||||
// Second worker of same model uses existing policy
|
||||
let policy2 = registry.on_worker_added("llama-3", Some("round_robin"));
|
||||
assert_eq!(policy2.name(), "cache_aware"); // Ignores hint, uses existing
|
||||
|
||||
// Different model can have different policy
|
||||
let policy3 = registry.on_worker_added("gpt-4", Some("random"));
|
||||
assert_eq!(policy3.name(), "random");
|
||||
|
||||
// Check mappings
|
||||
let mappings = registry.get_all_mappings();
|
||||
assert_eq!(mappings.get("llama-3").unwrap(), "cache_aware");
|
||||
assert_eq!(mappings.get("gpt-4").unwrap(), "random");
|
||||
|
||||
// Check worker counts
|
||||
let counts = registry.get_worker_counts();
|
||||
assert_eq!(*counts.get("llama-3").unwrap(), 2);
|
||||
assert_eq!(*counts.get("gpt-4").unwrap(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_policy_registry_cleanup() {
|
||||
let registry = PolicyRegistry::new(PolicyConfig::RoundRobin);
|
||||
|
||||
// Add workers
|
||||
registry.on_worker_added("llama-3", Some("cache_aware"));
|
||||
registry.on_worker_added("llama-3", None);
|
||||
assert_eq!(registry.get_worker_counts().get("llama-3"), Some(&2));
|
||||
|
||||
// Remove one worker - policy should remain
|
||||
registry.on_worker_removed("llama-3");
|
||||
assert!(registry.get_policy("llama-3").is_some());
|
||||
assert_eq!(registry.get_worker_counts().get("llama-3"), Some(&1));
|
||||
|
||||
// Remove last worker - policy should be cleaned up
|
||||
registry.on_worker_removed("llama-3");
|
||||
assert!(registry.get_policy("llama-3").is_none());
|
||||
assert_eq!(registry.get_worker_counts().get("llama-3"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_default_policy() {
|
||||
let registry = PolicyRegistry::new(PolicyConfig::RoundRobin);
|
||||
|
||||
// No hint, no template - uses default
|
||||
let policy = registry.on_worker_added("unknown-model", None);
|
||||
assert_eq!(policy.name(), "round_robin");
|
||||
|
||||
// Get default directly
|
||||
let default = registry.get_default_policy();
|
||||
assert_eq!(default.name(), "round_robin");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user