diff --git a/sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs b/sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs index c330d295b..635cd309d 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs @@ -32,7 +32,10 @@ pub struct ServerInfo { pub model_id: Option, pub model_path: Option, pub served_model_name: Option, + pub tp_size: Option, pub dp_size: Option, + pub load_balance_method: Option, + pub disaggregation_mode: Option, pub version: Option, pub max_batch_size: Option, pub max_total_tokens: Option, @@ -242,6 +245,18 @@ impl StepExecutor for DiscoverMetadataStep { { labels.insert("served_model_name".to_string(), served_model_name); } + if let Some(tp_size) = server_info.tp_size { + labels.insert("tp_size".to_string(), tp_size.to_string()); + } + if let Some(dp_size) = server_info.dp_size { + labels.insert("dp_size".to_string(), dp_size.to_string()); + } + if let Some(load_balance_method) = server_info.load_balance_method { + labels.insert("load_balance_method".to_string(), load_balance_method); + } + if let Some(disaggregation_mode) = server_info.disaggregation_mode { + labels.insert("disaggregation_mode".to_string(), disaggregation_mode); + } } // Fetch from /model_info for model-related metadata diff --git a/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs b/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs index e2733066f..62ae9cbd5 100644 --- a/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs +++ b/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs @@ -3,7 +3,7 @@ use std::{collections::HashMap, sync::Arc}; use async_trait::async_trait; -use tracing::debug; +use tracing::{debug, warn}; use crate::{ app_context::AppContext, @@ -17,6 +17,70 @@ use crate::{ /// external workers (different models per worker). pub struct UpdatePoliciesStep; +impl UpdatePoliciesStep { + /// Check for conflicts between prefill and decode worker configurations for a model. + fn check_worker_conflicts(&self, model_id: &str, workers: &[Arc]) { + let prefill_workers: Vec<_> = workers + .iter() + .filter(|w| { + w.metadata() + .labels + .get("disaggregation_mode") + .map(|s| s.as_str()) + == Some("prefill") + }) + .collect(); + + let decode_workers: Vec<_> = workers + .iter() + .filter(|w| { + w.metadata() + .labels + .get("disaggregation_mode") + .map(|s| s.as_str()) + == Some("decode") + }) + .collect(); + + if prefill_workers.is_empty() || decode_workers.is_empty() { + return; + } + + // Compare configurations of prefill vs decode workers + if let (Some(pw), Some(dw)) = (prefill_workers.first(), decode_workers.first()) { + let pl = &pw.metadata().labels; + let dl = &dw.metadata().labels; + + // Define keys to check for equality + let keys_to_check = ["tp_size", "dp_size", "load_balance_method"]; + + for key in keys_to_check { + let p_val = pl.get(key); + let d_val = dl.get(key); + if p_val != d_val { + warn!( + "Model {} has conflicting {}: prefill={:?}, decode={:?}", + model_id, key, p_val, d_val + ); + } + } + + // Specific check for Data-Parallel consistency + if let Some(dp_size) = pl.get("dp_size").and_then(|s| s.parse::().ok()) { + if dp_size > 1 { + let plb = pl.get("load_balance_method").map(|s| s.as_str()); + if plb != Some("follow_bootstrap_room") { + warn!( + "Model {} has dp_size > 1 but load_balance_method is not 'follow_bootstrap_room' on prefill workers. This may cause rank mismatch in disaggregated mode.", + model_id + ); + } + } + } + } + } +} + #[async_trait] impl StepExecutor for UpdatePoliciesStep { async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { @@ -39,6 +103,9 @@ impl StepExecutor for UpdatePoliciesStep { // Initialize cache-aware policy if configured let all_workers = app_context.worker_registry.get_by_model(&model_id); + + // Check for configuration conflicts between prefill and decode + self.check_worker_conflicts(&model_id, &all_workers); if let Some(policy) = app_context.policy_registry.get_policy(&model_id) { if policy.name() == "cache_aware" { app_context