[model-gateway]: move PD configuration conflict checks to model gateway (#16088)
This commit is contained in:
@@ -32,7 +32,10 @@ pub struct ServerInfo {
|
|||||||
pub model_id: Option<String>,
|
pub model_id: Option<String>,
|
||||||
pub model_path: Option<String>,
|
pub model_path: Option<String>,
|
||||||
pub served_model_name: Option<String>,
|
pub served_model_name: Option<String>,
|
||||||
|
pub tp_size: Option<usize>,
|
||||||
pub dp_size: Option<usize>,
|
pub dp_size: Option<usize>,
|
||||||
|
pub load_balance_method: Option<String>,
|
||||||
|
pub disaggregation_mode: Option<String>,
|
||||||
pub version: Option<String>,
|
pub version: Option<String>,
|
||||||
pub max_batch_size: Option<usize>,
|
pub max_batch_size: Option<usize>,
|
||||||
pub max_total_tokens: Option<usize>,
|
pub max_total_tokens: Option<usize>,
|
||||||
@@ -242,6 +245,18 @@ impl StepExecutor for DiscoverMetadataStep {
|
|||||||
{
|
{
|
||||||
labels.insert("served_model_name".to_string(), served_model_name);
|
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
|
// Fetch from /model_info for model-related metadata
|
||||||
|
|||||||
@@ -3,7 +3,7 @@
|
|||||||
use std::{collections::HashMap, sync::Arc};
|
use std::{collections::HashMap, sync::Arc};
|
||||||
|
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use tracing::debug;
|
use tracing::{debug, warn};
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
app_context::AppContext,
|
app_context::AppContext,
|
||||||
@@ -17,6 +17,70 @@ use crate::{
|
|||||||
/// external workers (different models per worker).
|
/// external workers (different models per worker).
|
||||||
pub struct UpdatePoliciesStep;
|
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<dyn Worker>]) {
|
||||||
|
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::<usize>().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]
|
#[async_trait]
|
||||||
impl StepExecutor for UpdatePoliciesStep {
|
impl StepExecutor for UpdatePoliciesStep {
|
||||||
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
|
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
|
||||||
@@ -39,6 +103,9 @@ impl StepExecutor for UpdatePoliciesStep {
|
|||||||
|
|
||||||
// Initialize cache-aware policy if configured
|
// Initialize cache-aware policy if configured
|
||||||
let all_workers = app_context.worker_registry.get_by_model(&model_id);
|
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 let Some(policy) = app_context.policy_registry.get_policy(&model_id) {
|
||||||
if policy.name() == "cache_aware" {
|
if policy.name() == "cache_aware" {
|
||||||
app_context
|
app_context
|
||||||
|
|||||||
Reference in New Issue
Block a user