diff --git a/sgl-router/src/core/workflow/steps/worker_registration.rs b/sgl-router/src/core/workflow/steps/worker_registration.rs index f2b7a959e..a214027d1 100644 --- a/sgl-router/src/core/workflow/steps/worker_registration.rs +++ b/sgl-router/src/core/workflow/steps/worker_registration.rs @@ -44,6 +44,7 @@ struct ServerInfo { #[serde(alias = "model")] model_id: Option, model_path: Option, + served_model_name: Option, dp_size: Option, version: Option, max_batch_size: Option, @@ -105,6 +106,8 @@ async fn get_dp_info(url: &str, api_key: Option<&str>) -> Result let model_id = info .model_id + .filter(|s| !s.is_empty()) + .or(info.served_model_name.filter(|s| !s.is_empty())) .or_else(|| { info.model_path .and_then(|path| path.split('/').next_back().map(|s| s.to_string())) @@ -318,20 +321,23 @@ impl StepExecutor for DiscoverMetadataStep { ); let discovered_labels = match connection_mode.as_ref() { - ConnectionMode::Http => { - match get_server_info(&config.url, config.api_key.as_deref()).await { - Ok(server_info) => { - let mut labels = HashMap::new(); - if let Some(model_path) = server_info.model_path { - if !model_path.is_empty() { - labels.insert("model_path".to_string(), model_path); - } + ConnectionMode::Http => match get_server_info(&config.url, config.api_key.as_deref()).await { + Ok(server_info) => { + let mut labels = HashMap::new(); + if let Some(model_path) = server_info.model_path { + if !model_path.is_empty() { + labels.insert("model_path".to_string(), model_path); } - Ok(labels) } - Err(e) => Err(e), + if let Some(served_model_name) = server_info.served_model_name { + if !served_model_name.is_empty() { + labels.insert("served_model_name".to_string(), served_model_name); + } + } + Ok(labels) } - } + Err(e) => Err(e), + }, ConnectionMode::Grpc { .. } => fetch_grpc_metadata(&config.url).await, } .unwrap_or_else(|e| {