From 5dad1ff1b6dd534c893c55760f655630ef225f9e Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Tue, 2 Dec 2025 16:47:06 -0800 Subject: [PATCH] [model-gateway] add workflow for external model providers (#14323) --- sgl-router/src/core/job_queue.rs | 76 ++- sgl-router/src/core/workflow/mod.rs | 4 +- .../steps/external_worker_registration.rs | 586 ++++++++++++++++++ sgl-router/src/core/workflow/steps/mod.rs | 14 +- .../workflow/steps/worker_registration.rs | 161 ++--- sgl-router/src/server.rs | 6 +- 6 files changed, 719 insertions(+), 128 deletions(-) create mode 100644 sgl-router/src/core/workflow/steps/external_worker_registration.rs diff --git a/sgl-router/src/core/job_queue.rs b/sgl-router/src/core/job_queue.rs index bd25dfbc9..b47fe3919 100644 --- a/sgl-router/src/core/job_queue.rs +++ b/sgl-router/src/core/job_queue.rs @@ -371,9 +371,71 @@ impl JobQueue { prefill_workers.chain(decode_workers).collect() } - RoutingMode::OpenAI { .. } => { - info!("OpenAI mode: no workers to initialize"); - return Ok("OpenAI mode: no workers to initialize".to_string()); + RoutingMode::OpenAI { worker_urls } => { + // OpenAI mode: submit AddWorker jobs with runtime: "external" + // The external_worker_registration workflow handles model discovery + let api_key = router_config.api_key.clone(); + let mut submitted_count = 0; + + for url in worker_urls { + let url_for_error = url.clone(); + let config = WorkerConfigRequest { + url: url.clone(), + api_key: api_key.clone(), + worker_type: Some("regular".to_string()), + labels: HashMap::new(), + model_id: None, + priority: None, + cost: None, + runtime: Some("external".to_string()), + tokenizer_path: None, + reasoning_parser: None, + tool_parser: None, + chat_template: None, + bootstrap_port: None, + health_check_timeout_secs: router_config.health_check.timeout_secs, + health_check_interval_secs: router_config + .health_check + .check_interval_secs, + health_success_threshold: router_config + .health_check + .success_threshold, + health_failure_threshold: router_config + .health_check + .failure_threshold, + max_connection_attempts: router_config + .health_check + .success_threshold + * 10, + dp_aware: false, + }; + + let job = Job::AddWorker { + config: Box::new(config), + }; + + if let Some(queue) = context.worker_job_queue.get() { + queue.submit(job).await.map_err(|e| { + format!( + "Failed to submit AddWorker job for external endpoint {}: {}", + url_for_error, e + ) + })?; + submitted_count += 1; + } else { + return Err("JobQueue not available".to_string()); + } + } + + if submitted_count == 0 { + info!("OpenAI mode: no worker URLs provided"); + return Ok("OpenAI mode: no worker URLs to initialize".to_string()); + } + + return Ok(format!( + "Submitted {} AddWorker jobs for external endpoints", + submitted_count + )); } }; @@ -497,8 +559,14 @@ impl JobQueue { workflow_context.set("worker_config", config.clone()); workflow_context.set_arc("app_context", Arc::clone(context)); + // Select workflow based on runtime field + let workflow_id = match config.runtime.as_deref() { + Some("external") => WorkflowId::new("external_worker_registration"), + _ => WorkflowId::new("worker_registration"), + }; + engine - .start_workflow(WorkflowId::new("worker_registration"), workflow_context) + .start_workflow(workflow_id, workflow_context) .await .map_err(|e| format!("Failed to start worker registration workflow: {:?}", e)) } diff --git a/sgl-router/src/core/workflow/mod.rs b/sgl-router/src/core/workflow/mod.rs index 320ffe053..341ebea38 100644 --- a/sgl-router/src/core/workflow/mod.rs +++ b/sgl-router/src/core/workflow/mod.rs @@ -15,7 +15,7 @@ pub use event::{EventBus, EventSubscriber, LoggingSubscriber, WorkflowEvent}; pub use executor::{FunctionStep, StepExecutor}; pub use state::WorkflowStateStore; pub use steps::{ - create_mcp_registration_workflow, create_worker_registration_workflow, - create_worker_removal_workflow, + create_external_worker_registration_workflow, create_mcp_registration_workflow, + create_worker_registration_workflow, create_worker_removal_workflow, }; pub use types::*; diff --git a/sgl-router/src/core/workflow/steps/external_worker_registration.rs b/sgl-router/src/core/workflow/steps/external_worker_registration.rs new file mode 100644 index 000000000..aaefac8c1 --- /dev/null +++ b/sgl-router/src/core/workflow/steps/external_worker_registration.rs @@ -0,0 +1,586 @@ +//! External worker registration workflow steps +//! +//! This workflow handles registration of external API endpoints (OpenAI, xAI, Anthropic, etc.) +//! +//! Key features: +//! - Fetches models from /v1/models endpoint +//! - Groups dated model variants under base model names (e.g., gpt-4o, gpt-4o-2024-08-06) +//! - Infers ModelType from model ID patterns (LLM, embedding, image gen, audio, etc.) +//! +//! Workflow order: +//! 1. DiscoverModels - Fetch available models from /v1/models endpoint +//! 2. CreateExternalWorkers - Build worker objects for each discovered model +//! 3. RegisterWorkers - Register workers in registry +//! 4. UpdatePolicies - Update policy registry +//! 5. ActivateWorkers - Mark workers as healthy + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use async_trait::async_trait; +use once_cell::sync::Lazy; +use regex::Regex; +use reqwest::Client; +use serde::Deserialize; +use tracing::{debug, info}; + +use crate::{ + app_context::AppContext, + core::{ + model_card::{ModelCard, ProviderType}, + model_type::ModelType, + workflow::*, + BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, HealthConfig, RuntimeType, + Worker, WorkerType, + }, + protocols::worker_spec::WorkerConfigRequest, +}; + +// HTTP client for API calls +static HTTP_CLIENT: Lazy = Lazy::new(|| { + Client::builder() + .timeout(Duration::from_secs(30)) + .build() + .expect("Failed to create HTTP client") +}); + +// Regex to strip date suffix: -YYYY-MM-DD or -YYYY-MM +static DATE_SUFFIX_PATTERN: Lazy = + Lazy::new(|| Regex::new(r"-\d{4}-\d{2}(-\d{2})?$").expect("Invalid date regex")); + +// ============================================================================ +// Model Discovery Types +// ============================================================================ + +/// OpenAI /v1/models response format +#[derive(Debug, Clone, Deserialize)] +pub struct ModelsResponse { + pub data: Vec, + #[serde(default)] + pub object: String, +} + +/// Individual model information from /v1/models +#[derive(Debug, Clone, Deserialize)] +pub struct ModelInfo { + pub id: String, + #[serde(default)] + pub object: String, + #[serde(default)] + pub created: Option, + #[serde(default)] + pub owned_by: Option, +} + +// ============================================================================ +// Model Discovery Functions +// ============================================================================ + +/// Group models by base name (stripping date suffixes) and create ModelCards with aliases. +/// +/// # Example +/// Input: `["gpt-4o", "gpt-4o-2024-05-13", "gpt-4o-2024-08-06", "gpt-4o-2024-11-20"]` +/// Output: `ModelCard { id: "gpt-4o", aliases: ["gpt-4o-2024-05-13", "gpt-4o-2024-08-06", "gpt-4o-2024-11-20"] }` +pub fn group_models_into_cards(models: Vec) -> Vec { + // Group model IDs by base name (with date stripped) + let mut groups: HashMap> = HashMap::new(); + for model in &models { + let base = DATE_SUFFIX_PATTERN.replace(&model.id, "").to_string(); + groups.entry(base).or_default().push(model.id.clone()); + } + + // Create ModelCard for each group + groups + .into_values() + .map(|mut variants| { + // Sort: shortest first (base name), then alphabetically + variants.sort_by(|a, b| a.len().cmp(&b.len()).then_with(|| a.cmp(b))); + + let primary_id = variants.remove(0); // shortest = primary ID + let aliases = variants; // rest = aliases + + let model_type = infer_model_type_from_id(&primary_id); + let provider = infer_provider_from_id(&primary_id); + + let mut card = ModelCard::new(&primary_id) + .with_aliases(aliases) + .with_model_type(model_type); + + if let Some(p) = provider { + card = card.with_provider(p); + } + + card + }) + .collect() +} + +/// Infer ModelType from model ID string. +/// +/// This function analyzes the model ID to determine what type of model it is +/// (LLM, embedding, image generation, audio, etc.). +pub fn infer_model_type_from_id(id: &str) -> ModelType { + let id_lower = id.to_lowercase(); + + // Embedding models + if id_lower.contains("embed") || id_lower.contains("ada-002") { + return ModelType::EMBED_MODEL; + } + + // Rerank models + if id_lower.contains("rerank") { + return ModelType::RERANK_MODEL; + } + + // Image generation models (DALL-E, Sora, gpt-image) + if id_lower.starts_with("dall-e") + || id_lower.starts_with("sora") + || (id_lower.contains("image") && !id_lower.contains("vision")) + { + return ModelType::IMAGE_MODEL; + } + + // Audio models (TTS, Whisper, realtime, audio) + if id_lower.starts_with("tts") + || id_lower.starts_with("whisper") + || id_lower.contains("audio") + || id_lower.contains("realtime") + || id_lower.contains("transcribe") + { + return ModelType::AUDIO_MODEL; + } + + // Moderation models + if id_lower.contains("moderation") { + return ModelType::MODERATION_MODEL; + } + + // Vision LLM (models with vision capability) + // gpt-4o, gpt-4-vision, etc. + if id_lower.contains("vision") || id_lower.contains("4o") { + return ModelType::VISION_LLM; + } + + // Reasoning models (o1, o3, etc.) + if id_lower.starts_with("o1") || id_lower.starts_with("o3") { + return ModelType::REASONING_LLM; + } + + // Default to standard LLM + ModelType::LLM +} + +/// Infer provider type from model ID string. +/// +/// Returns `None` for models that don't clearly indicate a provider. +fn infer_provider_from_id(id: &str) -> Option { + let id_lower = id.to_lowercase(); + + // OpenAI models + if id_lower.starts_with("gpt") + || id_lower.starts_with("o1") + || id_lower.starts_with("o3") + || id_lower.starts_with("dall-e") + || id_lower.starts_with("whisper") + || id_lower.starts_with("tts") + || id_lower.starts_with("text-embedding") + || id_lower.starts_with("babbage") + || id_lower.starts_with("davinci") + || id_lower.contains("omni") + { + return Some(ProviderType::OpenAI); + } + + // xAI/Grok models + if id_lower.starts_with("grok") { + return Some(ProviderType::XAI); + } + + // Anthropic Claude models + if id_lower.starts_with("claude") { + return Some(ProviderType::Anthropic); + } + + // Google Gemini models + if id_lower.starts_with("gemini") { + return Some(ProviderType::Gemini); + } + + None +} + +// ============================================================================ +// Workflow Steps +// ============================================================================ + +/// Step 1: Discover models from external /v1/models endpoint +pub struct DiscoverModelsStep; + +#[async_trait] +impl StepExecutor for DiscoverModelsStep { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + let config: Arc = context + .get("worker_config") + .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; + + debug!("Discovering models from external endpoint {}", config.url); + + let model_cards = fetch_models(&config.url, config.api_key.as_deref()) + .await + .map_err(|e| WorkflowError::StepFailed { + step_id: StepId::new("discover_models"), + message: format!("Failed to discover models from {}: {}", config.url, e), + })?; + + if model_cards.is_empty() { + return Err(WorkflowError::StepFailed { + step_id: StepId::new("discover_models"), + message: format!("No models discovered from {}", config.url), + }); + } + + info!( + "Discovered {} models from {}: {:?}", + model_cards.len(), + config.url, + model_cards.iter().map(|c| &c.id).collect::>() + ); + + context.set("model_cards", model_cards); + Ok(StepResult::Success) + } + + fn is_retryable(&self, _error: &WorkflowError) -> bool { + true // Network issues are retryable + } +} + +/// Fetch models from /v1/models endpoint +async fn fetch_models(url: &str, api_key: Option<&str>) -> Result, String> { + let base_url = url.trim_end_matches('/'); + let models_url = format!("{}/v1/models", base_url); + + let mut req = HTTP_CLIENT.get(&models_url); + if let Some(key) = api_key { + req = req.bearer_auth(key); + } + + let response = req + .send() + .await + .map_err(|e| format!("Failed to connect to {}: {}", models_url, e))?; + + if !response.status().is_success() { + return Err(format!( + "Server returned status {} from {}", + response.status(), + models_url + )); + } + + let models_response: ModelsResponse = response + .json() + .await + .map_err(|e| format!("Failed to parse models response: {}", e))?; + + debug!( + "Fetched {} raw models from {}", + models_response.data.len(), + url + ); + + // Group models into cards (e.g., gpt-4o-2024-08-06 → gpt-4o with aliases) + let model_cards = group_models_into_cards(models_response.data); + + debug!( + "Grouped into {} model cards with aliases", + model_cards.len() + ); + + Ok(model_cards) +} + +/// Step 2: Create worker objects for each discovered model +pub struct CreateExternalWorkersStep; + +#[async_trait] +impl StepExecutor for CreateExternalWorkersStep { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + let config: Arc = context + .get("worker_config") + .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; + let app_context: Arc = context + .get("app_context") + .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; + let model_cards: Arc> = context + .get("model_cards") + .ok_or_else(|| WorkflowError::ContextValueNotFound("model_cards".to_string()))?; + + debug!( + "Creating {} external workers for {}", + model_cards.len(), + config.url + ); + + // Build configs from router settings + let circuit_breaker_config = { + let cfg = app_context.router_config.effective_circuit_breaker_config(); + CircuitBreakerConfig { + failure_threshold: cfg.failure_threshold, + success_threshold: cfg.success_threshold, + timeout_duration: Duration::from_secs(cfg.timeout_duration_secs), + window_duration: Duration::from_secs(cfg.window_duration_secs), + } + }; + + let health_config = { + let cfg = &app_context.router_config.health_check; + HealthConfig { + timeout_secs: cfg.timeout_secs, + check_interval_secs: cfg.check_interval_secs, + endpoint: cfg.endpoint.clone(), + failure_threshold: cfg.failure_threshold, + success_threshold: cfg.success_threshold, + } + }; + + // Build labels + let mut labels = config.labels.clone(); + if let Some(priority) = config.priority { + labels.insert("priority".to_string(), priority.to_string()); + } + if let Some(cost) = config.cost { + labels.insert("cost".to_string(), cost.to_string()); + } + + // Normalize URL (ensure https:// for external APIs) + let normalized_url = normalize_external_url(&config.url); + + // Create a worker for each model + let mut workers = Vec::new(); + for model_card in model_cards.iter() { + let mut builder = BasicWorkerBuilder::new(normalized_url.clone()) + .model(model_card.clone()) + .worker_type(WorkerType::Regular) + .connection_mode(ConnectionMode::Http) + .runtime_type(RuntimeType::External) + .circuit_breaker_config(circuit_breaker_config.clone()) + .health_config(health_config.clone()); + + if let Some(ref api_key) = config.api_key { + builder = builder.api_key(api_key.clone()); + } + + if !labels.is_empty() { + builder = builder.labels(labels.clone()); + } + + let worker = Arc::new(builder.build()) as Arc; + worker.set_healthy(false); + + debug!( + "Created external worker for model {} at {}", + model_card.id, normalized_url + ); + + workers.push(worker); + } + + info!( + "Created {} external workers from {}", + workers.len(), + config.url + ); + + context.set("workers", workers); + context.set("labels", labels); + Ok(StepResult::Success) + } + + fn is_retryable(&self, _error: &WorkflowError) -> bool { + false // Creation failures are config issues + } +} + +/// Normalize URL for external APIs (ensure https://) +fn normalize_external_url(url: &str) -> String { + if url.starts_with("http://") || url.starts_with("https://") { + url.to_string() + } else { + format!("https://{}", url) + } +} + +/// Step 3: Register workers in registry +pub struct RegisterExternalWorkersStep; + +#[async_trait] +impl StepExecutor for RegisterExternalWorkersStep { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + let config: Arc = context + .get("worker_config") + .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; + let app_context: Arc = context + .get("app_context") + .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; + let workers: Arc>> = context + .get("workers") + .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?; + + let mut worker_ids = Vec::new(); + for worker in workers.iter() { + let worker_id = app_context.worker_registry.register(Arc::clone(worker)); + worker_ids.push(worker_id.clone()); + debug!( + "Registered external worker {} (model: {}) with ID {:?}", + config.url, + worker.model_id(), + worker_id + ); + } + + context.set("worker_ids", worker_ids); + Ok(StepResult::Success) + } + + fn is_retryable(&self, _error: &WorkflowError) -> bool { + false + } +} + +/// Step 4: Update policy registry +pub struct UpdateExternalPoliciesStep; + +#[async_trait] +impl StepExecutor for UpdateExternalPoliciesStep { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + let config: Arc = context + .get("worker_config") + .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; + let labels: Arc> = context + .get("labels") + .ok_or_else(|| WorkflowError::ContextValueNotFound("labels".to_string()))?; + let app_context: Arc = context + .get("app_context") + .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; + let workers: Arc>> = context + .get("workers") + .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?; + + let policy_hint = labels.get("policy").map(|s| s.as_str()); + + // Each external worker has a different model_id + for worker in workers.iter() { + let model_id = worker.model_id().to_string(); + app_context + .policy_registry + .on_worker_added(&model_id, policy_hint); + } + + debug!( + "Updated policies for {} external workers from {}", + workers.len(), + config.url + ); + + Ok(StepResult::Success) + } + + fn is_retryable(&self, _error: &WorkflowError) -> bool { + false + } +} + +/// Step 5: Activate workers by marking them healthy +pub struct ActivateExternalWorkersStep; + +#[async_trait] +impl StepExecutor for ActivateExternalWorkersStep { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + let config: Arc = context + .get("worker_config") + .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; + let workers: Arc>> = context + .get("workers") + .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?; + + for worker in workers.iter() { + worker.set_healthy(true); + } + + info!( + "Activated {} external workers from {} (marked as healthy)", + workers.len(), + config.url + ); + + Ok(StepResult::Success) + } + + fn is_retryable(&self, _error: &WorkflowError) -> bool { + false + } +} + +// ============================================================================ +// Workflow Definition +// ============================================================================ + +/// Create external worker registration workflow definition +pub fn create_external_worker_registration_workflow() -> WorkflowDefinition { + WorkflowDefinition::new( + "external_worker_registration", + "External Worker Registration", + ) + .add_step( + StepDefinition::new( + "discover_models", + "Discover Models", + Arc::new(DiscoverModelsStep), + ) + .with_retry(RetryPolicy { + max_attempts: 3, + backoff: BackoffStrategy::Exponential { + base: Duration::from_secs(1), + max: Duration::from_secs(10), + }, + }) + .with_timeout(Duration::from_secs(30)) + .with_failure_action(FailureAction::FailWorkflow), + ) + .add_step( + StepDefinition::new( + "create_workers", + "Create Workers", + Arc::new(CreateExternalWorkersStep), + ) + .with_timeout(Duration::from_secs(5)) + .with_failure_action(FailureAction::FailWorkflow), + ) + .add_step( + StepDefinition::new( + "register_workers", + "Register Workers", + Arc::new(RegisterExternalWorkersStep), + ) + .with_timeout(Duration::from_secs(5)) + .with_failure_action(FailureAction::FailWorkflow), + ) + .add_step( + StepDefinition::new( + "update_policies", + "Update Policies", + Arc::new(UpdateExternalPoliciesStep), + ) + .with_timeout(Duration::from_secs(5)) + .with_failure_action(FailureAction::ContinueNextStep), + ) + .add_step( + StepDefinition::new( + "activate_workers", + "Activate Workers", + Arc::new(ActivateExternalWorkersStep), + ) + .with_timeout(Duration::from_secs(5)) + .with_failure_action(FailureAction::FailWorkflow), + ) +} diff --git a/sgl-router/src/core/workflow/steps/mod.rs b/sgl-router/src/core/workflow/steps/mod.rs index 10b5c5f40..a3f67cc0d 100644 --- a/sgl-router/src/core/workflow/steps/mod.rs +++ b/sgl-router/src/core/workflow/steps/mod.rs @@ -1,22 +1,30 @@ //! Workflow step implementations //! //! This module contains concrete step implementations for various workflows: -//! - Worker registration and activation +//! - Local worker registration (SGLang, vLLM - gRPC or HTTP modes) +//! - External worker registration (OpenAI, xAI, Anthropic, etc. - HTTPS only) //! - Worker removal //! - MCP server registration -//! - Future: Tokenizer fetching, LoRA updates, etc. +pub mod external_worker_registration; pub mod mcp_registration; pub mod worker_registration; pub mod worker_removal; +pub use external_worker_registration::{ + create_external_worker_registration_workflow, group_models_into_cards, + infer_model_type_from_id, ActivateExternalWorkersStep, CreateExternalWorkersStep, + DiscoverModelsStep, ModelInfo, ModelsResponse, RegisterExternalWorkersStep, + UpdateExternalPoliciesStep, +}; pub use mcp_registration::{ create_mcp_registration_workflow, ConnectMcpServerStep, DiscoverMcpInventoryStep, McpServerConfigRequest, RegisterMcpServerStep, ValidateRegistrationStep, }; pub use worker_registration::{ create_worker_registration_workflow, ActivateWorkerStep, CreateWorkerStep, - DetectConnectionModeStep, DiscoverMetadataStep, RegisterWorkerStep, UpdatePoliciesStep, + DetectConnectionModeStep, DiscoverDPInfoStep, DiscoverMetadataStep, RegisterWorkerStep, + UpdatePoliciesStep, }; pub use worker_removal::{ create_worker_removal_workflow, FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep, diff --git a/sgl-router/src/core/workflow/steps/worker_registration.rs b/sgl-router/src/core/workflow/steps/worker_registration.rs index 7b380f072..c8b4a09ca 100644 --- a/sgl-router/src/core/workflow/steps/worker_registration.rs +++ b/sgl-router/src/core/workflow/steps/worker_registration.rs @@ -1,10 +1,11 @@ -//! Worker registration workflow steps +//! Local worker registration workflow steps //! -//! Each step is atomic and performs a single operation in the worker registration process. +//! This workflow handles registration of local inference workers (SGLang, vLLM). +//! For external API endpoints (OpenAI, xAI, etc.), see external_worker_registration.rs. //! //! Workflow order: -//! 1. DetectConnectionMode - Probe both HTTP and gRPC to determine connection mode -//! 2. DiscoverMetadata - Fetch metadata from the worker +//! 1. DetectConnectionMode - Probe HTTP and gRPC to determine connection mode +//! 2. DiscoverMetadata - Fetch metadata from /get_server_info or gRPC //! 3. DiscoverDPInfo - Fetch DP (Data Parallel) information (only for DP-aware workers) //! 4. CreateWorker - Build worker object(s) with merged config + metadata //! 5. RegisterWorker - Register worker(s) in registry @@ -38,7 +39,7 @@ static HTTP_CLIENT: Lazy = Lazy::new(|| { .expect("Failed to create HTTP client") }); -/// Server information returned from worker endpoints +/// Server information returned from /get_server_info endpoint #[derive(Debug, Clone, Deserialize, Serialize)] struct ServerInfo { #[serde(alias = "model")] @@ -60,7 +61,7 @@ pub struct DpInfo { pub model_id: String, } -/// Parse server info from JSON response using serde +/// Parse server info from JSON response fn parse_server_info(json: Value) -> Result { serde_json::from_value(json).map_err(|e| format!("Failed to parse server info: {}", e)) } @@ -117,7 +118,7 @@ async fn get_dp_info(url: &str, api_key: Option<&str>) -> Result Ok(DpInfo { dp_size, model_id }) } -/// Helper: Strip protocol prefix from URL +/// Strip protocol prefix from URL fn strip_protocol(url: &str) -> String { url.trim_start_matches("http://") .trim_start_matches("https://") @@ -125,24 +126,17 @@ fn strip_protocol(url: &str) -> String { .to_string() } -/// Helper: Try HTTP health check -/// -/// Uses the provided client (from app_context) which supports both HTTP and HTTPS. -/// For HTTPS URLs, the client's TLS configuration (mTLS, CA certs) is used. -/// For plain HTTP URLs, the client handles them normally without TLS overhead. +/// Try HTTP health check async fn try_http_health_check( url: &str, timeout_secs: u64, client: &Client, ) -> Result<(), String> { - // Preserve the protocol (http or https) from the original URL let is_https = url.starts_with("https://"); let protocol = if is_https { "https" } else { "http" }; let clean_url = strip_protocol(url); let health_url = format!("{}://{}/health", protocol, clean_url); - // Use the AppContext client for both HTTP and HTTPS - // The rustls backend handles both protocols correctly client .get(&health_url) .timeout(Duration::from_secs(timeout_secs)) @@ -154,7 +148,7 @@ async fn try_http_health_check( Ok(()) } -/// Helper: Perform gRPC health check with runtime type +/// Perform gRPC health check with runtime type async fn do_grpc_health_check( grpc_url: &str, timeout_secs: u64, @@ -175,10 +169,7 @@ async fn do_grpc_health_check( Ok(()) } -/// Helper: Try gRPC health check -/// -/// If runtime_type is specified, uses the appropriate client (SGLang or vLLM). -/// If not specified, tries SGLang first, then falls back to vLLM. +/// Try gRPC health check (tries SGLang first, then vLLM if not specified) async fn try_grpc_health_check( url: &str, timeout_secs: u64, @@ -193,25 +184,18 @@ async fn try_grpc_health_check( match runtime_type { Some(runtime) => do_grpc_health_check(&grpc_url, timeout_secs, runtime).await, None => { - // Runtime not specified: Try SGLang first, then vLLM as fallback + // Try SGLang first, then vLLM as fallback if let Ok(()) = do_grpc_health_check(&grpc_url, timeout_secs, "sglang").await { return Ok(()); } - - // Try vLLM as fallback do_grpc_health_check(&grpc_url, timeout_secs, "vllm") .await - .map_err(|e| { - format!( - "gRPC health check failed (tried both SGLang and vLLM): {}", - e - ) - }) + .map_err(|e| format!("gRPC failed (tried SGLang and vLLM): {}", e)) } } } -/// Fetch metadata from gRPC server with runtime type +/// Fetch metadata from gRPC server async fn do_fetch_grpc_metadata( grpc_url: &str, runtime_type: &str, @@ -228,11 +212,7 @@ async fn do_fetch_grpc_metadata( Ok(model_info.to_labels()) } -/// Helper: Fetch gRPC metadata -/// -/// If runtime_type is specified, uses the appropriate client (SGLang or vLLM). -/// If not specified, tries SGLang first, then falls back to vLLM. -/// Returns (labels, detected_runtime_type) +/// Fetch gRPC metadata (returns labels and detected runtime type) async fn fetch_grpc_metadata( url: &str, runtime_type: Option<&str>, @@ -249,26 +229,19 @@ async fn fetch_grpc_metadata( Ok((labels, runtime.to_string())) } None => { - // Runtime not specified: Try SGLang first, then vLLM as fallback + // Try SGLang first, then vLLM as fallback if let Ok(labels) = do_fetch_grpc_metadata(&grpc_url, "sglang").await { return Ok((labels, "sglang".to_string())); } - - // Try vLLM as fallback let labels = do_fetch_grpc_metadata(&grpc_url, "vllm") .await - .map_err(|e| { - format!( - "Failed to fetch gRPC metadata (tried both SGLang and vLLM): {}", - e - ) - })?; + .map_err(|e| format!("gRPC metadata failed (tried SGLang and vLLM): {}", e))?; Ok((labels, "vllm".to_string())) } } } -/// Step 1: Detect connection mode by probing both HTTP and gRPC +/// Step 1: Detect connection mode by probing HTTP and gRPC pub struct DetectConnectionModeStep; #[async_trait] @@ -286,12 +259,12 @@ impl StepExecutor for DetectConnectionModeStep { config.url, config.health_check_timeout_secs, config.max_connection_attempts ); - // Try both protocols in parallel using configured timeout - // Use the AppContext client which has TLS configuration (CA certs, client identity) + // Try both protocols in parallel let url = config.url.clone(); let timeout = config.health_check_timeout_secs; let client = &app_context.client; let runtime_type = config.runtime.as_deref(); + let (http_result, grpc_result) = tokio::join!( try_http_health_check(&url, timeout, client), try_grpc_health_check(&url, timeout, runtime_type) @@ -317,14 +290,12 @@ impl StepExecutor for DetectConnectionModeStep { } }; - // Store connection mode in context context.set("connection_mode", connection_mode); - Ok(StepResult::Success) } fn is_retryable(&self, _error: &WorkflowError) -> bool { - true // Connection issues are retryable + true } } @@ -359,7 +330,6 @@ impl StepExecutor for DiscoverMetadataStep { { labels.insert("served_model_name".to_string(), served_model_name); } - Ok((labels, None)) } Err(e) => Err(e), @@ -383,7 +353,6 @@ impl StepExecutor for DiscoverMetadataStep { config.url ); - // Store discovered labels and detected runtime in context context.set("discovered_labels", discovered_labels); if let Some(runtime) = detected_runtime { debug!("Detected runtime type: {}", runtime); @@ -394,7 +363,7 @@ impl StepExecutor for DiscoverMetadataStep { } fn is_retryable(&self, _error: &WorkflowError) -> bool { - true // Metadata discovery failures are retryable + true } } @@ -408,7 +377,6 @@ impl StepExecutor for DiscoverDPInfoStep { .get("worker_config") .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; - // Skip DP discovery if not DP-aware if !config.dp_aware { debug!( "Worker {} is not DP-aware, skipping DP discovery", @@ -419,7 +387,6 @@ impl StepExecutor for DiscoverDPInfoStep { debug!("Discovering DP info for {} (DP-aware)", config.url); - // Get DP info from worker let dp_info = get_dp_info(&config.url, config.api_key.as_deref()) .await .map_err(|e| WorkflowError::StepFailed { @@ -432,14 +399,12 @@ impl StepExecutor for DiscoverDPInfoStep { dp_info.dp_size, config.url, dp_info.model_id ); - // Store DP info in context context.set("dp_info", dp_info); - Ok(StepResult::Success) } fn is_retryable(&self, _error: &WorkflowError) -> bool { - true // DP info discovery failures are retryable + true } } @@ -474,7 +439,7 @@ impl StepExecutor for CreateWorkerStep { }); } - // Build labels from config (non-model-specific labels only) + // Build labels from config let mut config_labels = config.labels.clone(); if let Some(priority) = config.priority { config_labels.insert("priority".to_string(), priority.to_string()); @@ -501,7 +466,7 @@ impl StepExecutor for CreateWorkerStep { debug!("Using model_id: {}", model_id); } - // Create ModelCard with model-specific configuration + // Create ModelCard let model_card = { let mut card = ModelCard::new(&model_id); if let Some(ref tokenizer_path) = config.tokenizer_path { @@ -540,9 +505,8 @@ impl StepExecutor for CreateWorkerStep { }) .unwrap_or(WorkerType::Regular); - // Get detected runtime type (for gRPC workers) + // Get runtime type (for gRPC workers) let runtime_type = if matches!(connection_mode.as_ref(), ConnectionMode::Grpc { .. }) { - // Try to get detected runtime from context, fall back to config, or default to sglang if let Some(detected_runtime) = context.get::("detected_runtime_type") { match detected_runtime.as_str() { "vllm" => RuntimeType::Vllm, @@ -557,7 +521,7 @@ impl StepExecutor for CreateWorkerStep { RuntimeType::Sglang } } else { - RuntimeType::Sglang // Default for HTTP workers + RuntimeType::Sglang }; // Build circuit breaker config @@ -588,10 +552,8 @@ impl StepExecutor for CreateWorkerStep { || config.url.starts_with("https://") || config.url.starts_with("grpc://") { - // URL already has protocol, use as-is config.url.clone() } else { - // Bare IP:port format, add appropriate protocol based on detected mode match connection_mode.as_ref() { ConnectionMode::Http => format!("http://{}", config.url), ConnectionMode::Grpc { .. } => format!("grpc://{}", config.url), @@ -609,7 +571,6 @@ impl StepExecutor for CreateWorkerStep { // Handle DP-aware vs non-DP-aware workers if config.dp_aware { - // DP-aware path: Create multiple workers (one per rank) let dp_info: Arc = context .get("dp_info") .ok_or_else(|| WorkflowError::ContextValueNotFound("dp_info".to_string()))?; @@ -633,7 +594,6 @@ impl StepExecutor for CreateWorkerStep { if let Some(ref api_key) = config.api_key { builder = builder.api_key(api_key.clone()); } - if !final_labels.is_empty() { builder = builder.labels(final_labels.clone()); } @@ -651,13 +611,9 @@ impl StepExecutor for CreateWorkerStep { ); } - // Store workers (plural) and labels in context context.set("workers", workers); context.set("labels", final_labels); - - Ok(StepResult::Success) } else { - // Non-DP-aware path: Create single worker let mut builder = BasicWorkerBuilder::new(normalized_url.clone()) .model(model_card) .worker_type(worker_type) @@ -669,7 +625,6 @@ impl StepExecutor for CreateWorkerStep { if let Some(ref api_key) = config.api_key { builder = builder.api_key(api_key.clone()); } - if !final_labels.is_empty() { builder = builder.labels(final_labels.clone()); } @@ -684,16 +639,15 @@ impl StepExecutor for CreateWorkerStep { final_labels.len() ); - // Store worker (singular) and labels in context context.set("worker", worker); context.set("labels", final_labels); - - Ok(StepResult::Success) } + + Ok(StepResult::Success) } fn is_retryable(&self, _error: &WorkflowError) -> bool { - false // Worker creation failures are not retryable (likely config issues) + false } } @@ -710,9 +664,7 @@ impl StepExecutor for RegisterWorkerStep { .get("app_context") .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; - // Check if we have multiple workers (DP-aware) or single worker if config.dp_aware { - // DP-aware path: Register multiple workers let workers: Arc>> = context .get("workers") .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?; @@ -728,9 +680,7 @@ impl StepExecutor for RegisterWorkerStep { } context.set("worker_ids", worker_ids); - Ok(StepResult::Success) } else { - // Non-DP-aware path: Register single worker let worker: Arc> = context .get("worker") .ok_or_else(|| WorkflowError::ContextValueNotFound("worker".to_string()))?; @@ -738,16 +688,15 @@ impl StepExecutor for RegisterWorkerStep { let worker_id = app_context .worker_registry .register(Arc::clone(worker.as_ref())); - debug!("Registered worker {} with ID {:?}", config.url, worker_id); context.set("worker_id", worker_id); - - Ok(StepResult::Success) } + + Ok(StepResult::Success) } fn is_retryable(&self, _error: &WorkflowError) -> bool { - false // Registration failures are not retryable + false } } @@ -769,17 +718,13 @@ impl StepExecutor for UpdatePoliciesStep { let policy_hint = labels.get("policy").map(|s| s.as_str()); - // Check if we have multiple workers (DP-aware) or single worker if config.dp_aware { - // DP-aware path: Update policies for multiple workers let workers: Arc>> = context .get("workers") .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?; - // Get model_id from first worker (all DP workers have same model) let model_id = workers[0].model_id().to_string(); - // Notify policy registry for each worker for _ in 0..workers.len() { app_context .policy_registry @@ -803,14 +748,12 @@ impl StepExecutor for UpdatePoliciesStep { model_id ); } else { - // Non-DP-aware path: Update policy for single worker let worker: Arc> = context .get("worker") .ok_or_else(|| WorkflowError::ContextValueNotFound("worker".to_string()))?; let model_id = worker.model_id().to_string(); - // Notify policy registry app_context .policy_registry .on_worker_added(&model_id, policy_hint); @@ -824,6 +767,8 @@ impl StepExecutor for UpdatePoliciesStep { .init_cache_aware_policy(&model_id, &all_workers); } } + + // Initialize bucket policies for prefill workers let prefill_workers = app_context.worker_registry.get_prefill_workers(); let policy = app_context.policy_registry.get_prefill_policy(); if policy.name() == "bucket" { @@ -842,7 +787,7 @@ impl StepExecutor for UpdatePoliciesStep { } fn is_retryable(&self, _error: &WorkflowError) -> bool { - false // Policy update failures are not retryable + false } } @@ -856,9 +801,7 @@ impl StepExecutor for ActivateWorkerStep { .get("worker_config") .ok_or_else(|| WorkflowError::ContextValueNotFound("worker_config".to_string()))?; - // Check if we have multiple workers (DP-aware) or single worker if config.dp_aware { - // DP-aware path: Activate multiple workers let workers: Arc>> = context .get("workers") .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?; @@ -867,13 +810,12 @@ impl StepExecutor for ActivateWorkerStep { worker.set_healthy(true); } - debug!( - "Activated {} DP-aware workers {} (marked as healthy)", + info!( + "Activated {} DP-aware workers from {} (marked as healthy)", workers.len(), config.url ); } else { - // Non-DP-aware path: Activate single worker let worker: Arc> = context .get("worker") .ok_or_else(|| WorkflowError::ContextValueNotFound("worker".to_string()))?; @@ -887,31 +829,17 @@ impl StepExecutor for ActivateWorkerStep { } fn is_retryable(&self, _error: &WorkflowError) -> bool { - false // Activation is just setting a flag, not retryable + false } } -/// Create worker registration workflow definition -/// -/// Note: Actual health check timeouts and retry attempts are configured per-worker -/// via WorkerConfigRequest (populated from router config). The timeouts and retry -/// policies here serve as workflow-level bounds to prevent infinite waiting. -/// -/// # Arguments -/// * `router_config` - Router configuration containing health check settings +/// Create local worker registration workflow definition pub fn create_worker_registration_workflow( router_config: &crate::config::RouterConfig, ) -> WorkflowDefinition { - // Use startup timeout from config for worker registration - // This is separate from health_check.timeout_secs which is for individual HTTP requests let detect_timeout = Duration::from_secs(router_config.worker_startup_timeout_secs); - // Calculate max_attempts to match the startup_timeout - // With Linear backoff (increment 1s, max 5s): - // - Attempts 1-5: 0s, 1s, 2s, 3s, 4s = 10s total - // - Attempts 6+: 5s each - // max_attempts = 5 + (timeout_seconds - 10) / 5 - // Use 90% of timeout to leave buffer for actual connection attempts + // Calculate max_attempts based on timeout let timeout_secs = detect_timeout.as_secs() as f64; let effective_timeout = timeout_secs * 0.9; let max_attempts = if effective_timeout > 10.0 { @@ -934,7 +862,6 @@ pub fn create_worker_registration_workflow( max: Duration::from_secs(5), }, }) - // Workflow-level timeout uses configured health check timeout + buffer .with_timeout(detect_timeout) .with_failure_action(FailureAction::FailWorkflow), ) @@ -949,7 +876,7 @@ pub fn create_worker_registration_workflow( backoff: BackoffStrategy::Fixed(Duration::from_secs(1)), }) .with_timeout(Duration::from_secs(10)) - .with_failure_action(FailureAction::ContinueNextStep), // Metadata discovery is optional + .with_failure_action(FailureAction::ContinueNextStep), ) .add_step( StepDefinition::new( @@ -962,7 +889,7 @@ pub fn create_worker_registration_workflow( backoff: BackoffStrategy::Fixed(Duration::from_secs(1)), }) .with_timeout(Duration::from_secs(10)) - .with_failure_action(FailureAction::FailWorkflow), // DP info is required for DP-aware workers + .with_failure_action(FailureAction::FailWorkflow), ) .add_step( StepDefinition::new("create_worker", "Create Worker", Arc::new(CreateWorkerStep)) @@ -985,7 +912,7 @@ pub fn create_worker_registration_workflow( Arc::new(UpdatePoliciesStep), ) .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::ContinueNextStep), // Policy updates are optional + .with_failure_action(FailureAction::ContinueNextStep), ) .add_step( StepDefinition::new( diff --git a/sgl-router/src/server.rs b/sgl-router/src/server.rs index 7a6c6d3a0..79e6c87f4 100644 --- a/sgl-router/src/server.rs +++ b/sgl-router/src/server.rs @@ -24,8 +24,9 @@ use crate::{ core::{ worker_to_info, workflow::{ - create_mcp_registration_workflow, create_worker_registration_workflow, - create_worker_removal_workflow, LoggingSubscriber, WorkflowEngine, + create_external_worker_registration_workflow, create_mcp_registration_workflow, + create_worker_registration_workflow, create_worker_removal_workflow, LoggingSubscriber, + WorkflowEngine, }, Job, JobQueue, JobQueueConfig, WorkerManager, WorkerType, }, @@ -752,6 +753,7 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box