diff --git a/sgl-model-gateway/src/core/job_queue.rs b/sgl-model-gateway/src/core/job_queue.rs index fa2385661..06261a6cd 100644 --- a/sgl-model-gateway/src/core/job_queue.rs +++ b/sgl-model-gateway/src/core/job_queue.rs @@ -654,7 +654,7 @@ impl JobQueue { // 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"), + _ => WorkflowId::new("local_worker_registration"), }; engine diff --git a/sgl-model-gateway/src/core/steps/external_worker_registration.rs b/sgl-model-gateway/src/core/steps/external_worker_registration.rs deleted file mode 100644 index 640205843..000000000 --- a/sgl-model-gateway/src/core/steps/external_worker_registration.rs +++ /dev/null @@ -1,591 +0,0 @@ -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, - BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, HealthConfig, RuntimeType, - Worker, WorkerType, - }, - protocols::worker_spec::WorkerConfigRequest, - workflow::*, -}; - -// 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_or_err("worker_config")?; - - // If no API key is provided, skip model discovery and use wildcard mode. - if config.api_key.as_ref().is_none_or(|k| k.is_empty()) { - info!( - "No API key provided for {} - using wildcard mode (accepts any model). \ - User's Authorization header will be forwarded to backend.", - config.url - ); - context.set::>("model_cards", vec![]); - return Ok(StepResult::Success); - } - - 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_or_err("worker_config")?; - let app_context: Arc = context.get_or_err("app_context")?; - let model_cards: Arc> = context.get_or_err("model_cards")?; - - // 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); - - let mut workers = Vec::new(); - - // Handle wildcard mode: create a single worker with empty models list - if model_cards.is_empty() { - debug!("Creating wildcard worker (no models) for {}", config.url); - - let mut builder = BasicWorkerBuilder::new(normalized_url.clone()) - .models(vec![]) // Empty models = accepts any model - .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); - - info!( - "Created wildcard worker at {} (accepts any model, user auth forwarded)", - normalized_url - ); - - workers.push(worker); - } else { - debug!( - "Creating {} external workers for {}", - model_cards.len(), - config.url - ); - - // Create a worker for each model - 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_or_err("worker_config")?; - let app_context: Arc = context.get_or_err("app_context")?; - let workers: Arc>> = context.get_or_err("workers")?; - - 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_or_err("worker_config")?; - let labels: Arc> = context.get_or_err("labels")?; - let app_context: Arc = context.get_or_err("app_context")?; - let workers: Arc>> = context.get_or_err("workers")?; - - 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_or_err("worker_config")?; - let workers: Arc>> = context.get_or_err("workers")?; - - 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) - .depends_on(&["discover_models"]), - ) - .add_step( - StepDefinition::new( - "register_workers", - "Register Workers", - Arc::new(RegisterExternalWorkersStep), - ) - .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::FailWorkflow) - .depends_on(&["create_workers"]), - ) - .add_step( - StepDefinition::new( - "update_policies", - "Update Policies", - Arc::new(UpdateExternalPoliciesStep), - ) - .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::ContinueNextStep) - .depends_on(&["register_workers"]), - ) - .add_step( - StepDefinition::new( - "activate_workers", - "Activate Workers", - Arc::new(ActivateExternalWorkersStep), - ) - .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::FailWorkflow) - .depends_on(&["update_policies"]), - ) -} diff --git a/sgl-model-gateway/src/core/steps/mod.rs b/sgl-model-gateway/src/core/steps/mod.rs index 50b5cd4e0..e19165c48 100644 --- a/sgl-model-gateway/src/core/steps/mod.rs +++ b/sgl-model-gateway/src/core/steps/mod.rs @@ -1,26 +1,62 @@ //! Workflow step implementations //! //! This module contains concrete step implementations for various workflows: -//! - Local worker registration (SGLang, vLLM - gRPC or HTTP modes) -//! - External worker registration (OpenAI, xAI, Anthropic, etc. - HTTPS only) -//! - Worker removal +//! - Worker management (registration, removal, future: updates) //! - MCP server registration //! - WASM module registration and removal //! - Future: Tokenizer fetching, LoRA updates, etc. -pub mod external_worker_registration; pub mod mcp_registration; pub mod wasm_module_registration; pub mod wasm_module_removal; -pub mod worker_registration; -pub mod worker_removal; +pub mod worker; -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, +// Worker management (registration, removal) +#[allow(deprecated)] +pub use worker::create_external_worker_registration_workflow; +// Backward compatibility aliases +#[allow(deprecated)] +pub use worker::create_worker_registration_workflow; +pub use worker::{ + // Workflow builders + create_external_worker_workflow, + create_local_worker_workflow, + create_worker_removal_workflow, + // Utility functions + group_models_into_cards, + infer_model_type_from_id, + // Shared steps + ActivateWorkersStep, + // External registration steps + CreateExternalWorkersStep, + // Local registration steps + CreateLocalWorkerStep, + DetectConnectionModeStep, + DiscoverDPInfoStep, + DiscoverMetadataStep, + DiscoverModelsStep, + DpInfo, + // Removal steps + FindWorkersToRemoveStep, + ModelInfo, + ModelsResponse, + RegisterWorkersStep, + RemoveFromPolicyRegistryStep, + RemoveFromWorkerRegistryStep, + UpdatePoliciesStep, + UpdateRemainingPoliciesStep, + WorkerList, + WorkerRemovalRequest, }; + +// Legacy type aliases for backward compatibility +pub type ActivateWorkerStep = ActivateWorkersStep; +pub type RegisterWorkerStep = RegisterWorkersStep; +pub type CreateWorkerStep = CreateLocalWorkerStep; +pub type ActivateExternalWorkersStep = ActivateWorkersStep; +pub type RegisterExternalWorkersStep = RegisterWorkersStep; +pub type UpdateExternalPoliciesStep = UpdatePoliciesStep; + pub use mcp_registration::{ create_mcp_registration_workflow, ConnectMcpServerStep, DiscoverMcpInventoryStep, McpServerConfigRequest, RegisterMcpServerStep, ValidateRegistrationStep, @@ -34,12 +70,3 @@ pub use wasm_module_removal::{ create_wasm_module_removal_workflow, FindModuleToRemoveStep, RemoveModuleStep, WasmModuleRemovalRequest, }; -pub use worker_registration::{ - create_worker_registration_workflow, ActivateWorkerStep, CreateWorkerStep, - DetectConnectionModeStep, DiscoverDPInfoStep, DiscoverMetadataStep, RegisterWorkerStep, - UpdatePoliciesStep, -}; -pub use worker_removal::{ - create_worker_removal_workflow, FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep, - RemoveFromWorkerRegistryStep, UpdateRemainingPoliciesStep, WorkerRemovalRequest, -}; diff --git a/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs b/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs new file mode 100644 index 000000000..8ab6285d8 --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs @@ -0,0 +1,153 @@ +//! External worker creation step. + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use async_trait::async_trait; +use tracing::{debug, info}; + +use crate::{ + app_context::AppContext, + core::{ + model_card::ModelCard, BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, + HealthConfig, RuntimeType, Worker, WorkerType, + }, + protocols::worker_spec::WorkerConfigRequest, + workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, +}; + +/// 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 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_or_err("worker_config")?; + let app_context: Arc = context.get_or_err("app_context")?; + let model_cards: Arc> = context.get_or_err("model_cards")?; + + // 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 from config + let mut labels: HashMap = 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); + + let mut workers = Vec::new(); + + // Handle wildcard mode: create a single worker with empty models list + if model_cards.is_empty() { + debug!("Creating wildcard worker (no models) for {}", config.url); + + let mut builder = BasicWorkerBuilder::new(normalized_url.clone()) + .models(vec![]) // Empty models = accepts any model + .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); + + info!( + "Created wildcard worker at {} (accepts any model, user auth forwarded)", + normalized_url + ); + + workers.push(worker); + } else { + debug!( + "Creating {} external workers for {}", + model_cards.len(), + config.url + ); + + // Create a worker for each model + 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 + } +} diff --git a/sgl-model-gateway/src/core/steps/worker/external/discover_models.rs b/sgl-model-gateway/src/core/steps/worker/external/discover_models.rs new file mode 100644 index 000000000..066c01f49 --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/external/discover_models.rs @@ -0,0 +1,273 @@ +//! Model discovery step for external API endpoints. + +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::{ + core::{ + model_card::{ModelCard, ProviderType}, + model_type::ModelType, + }, + protocols::worker_spec::WorkerConfigRequest, + workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, +}; + +// 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")); + +/// 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, +} + +/// 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. +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 + 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 + 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 + if id_lower.contains("vision") || id_lower.contains("4o") { + return ModelType::VISION_LLM; + } + + // Reasoning models + 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. +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 +} + +/// 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 + ); + + let model_cards = group_models_into_cards(models_response.data); + + debug!( + "Grouped into {} model cards with aliases", + model_cards.len() + ); + + Ok(model_cards) +} + +/// 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_or_err("worker_config")?; + + // If no API key is provided, skip model discovery and use wildcard mode. + if config.api_key.as_ref().is_none_or(|k| k.is_empty()) { + info!( + "No API key provided for {} - using wildcard mode (accepts any model). \ + User's Authorization header will be forwarded to backend.", + config.url + ); + context.set::>("model_cards", vec![]); + return Ok(StepResult::Success); + } + + 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 + } +} diff --git a/sgl-model-gateway/src/core/steps/worker/external/mod.rs b/sgl-model-gateway/src/core/steps/worker/external/mod.rs new file mode 100644 index 000000000..99a9787f3 --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/external/mod.rs @@ -0,0 +1,104 @@ +//! External worker registration steps for OpenAI-compatible API endpoints. +//! +//! These steps handle the discovery and creation of workers that connect to +//! external API providers (OpenAI, Anthropic, etc.) via HTTP. + +mod create_workers; +mod discover_models; + +use std::{sync::Arc, time::Duration}; + +pub use create_workers::CreateExternalWorkersStep; +pub use discover_models::{ + group_models_into_cards, infer_model_type_from_id, DiscoverModelsStep, ModelInfo, + ModelsResponse, +}; + +use super::shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep}; +use crate::workflow::{ + BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition, +}; + +/// Create external worker registration workflow definition. +/// +/// DAG structure with parallel execution opportunities: +/// ```text +/// discover_models +/// │ +/// create_workers +/// │ +/// register_workers +/// │ +/// ┌────────────┴────────────┐ +/// │ │ +/// update_policies activate_workers +/// │ │ +/// └────────────┴────────────┘ +/// ``` +pub fn create_external_worker_workflow() -> WorkflowDefinition { + WorkflowDefinition::new( + "external_worker_registration", + "External Worker Registration", + ) + // Step 1: Discover models from /v1/models endpoint + .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), + ) + // Step 2: Create workers for each model + .add_step( + StepDefinition::new( + "create_workers", + "Create Workers", + Arc::new(CreateExternalWorkersStep), + ) + .with_timeout(Duration::from_secs(5)) + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["discover_models"]), + ) + // Step 3: Register workers (shared step) + .add_step( + StepDefinition::new( + "register_workers", + "Register Workers", + Arc::new(RegisterWorkersStep), + ) + .with_timeout(Duration::from_secs(5)) + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["create_workers"]), + ) + // Step 4a: Update policies (parallel with activation) + .add_step( + StepDefinition::new( + "update_policies", + "Update Policies", + Arc::new(UpdatePoliciesStep), + ) + .with_timeout(Duration::from_secs(5)) + .with_failure_action(FailureAction::ContinueNextStep) + .depends_on(&["register_workers"]), + ) + // Step 4b: Activate workers (parallel with policy update) + .add_step( + StepDefinition::new( + "activate_workers", + "Activate Workers", + Arc::new(ActivateWorkersStep), + ) + .with_timeout(Duration::from_secs(5)) + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["register_workers"]), + ) +} diff --git a/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs b/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs new file mode 100644 index 000000000..b4ec3095b --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs @@ -0,0 +1,340 @@ +//! Local worker creation step. + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use async_trait::async_trait; +use tracing::debug; + +use super::discover_dp::DpInfo; +use crate::{ + app_context::AppContext, + core::{ + BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, DPAwareWorkerBuilder, + HealthConfig, ModelCard, RuntimeType, Worker, WorkerType, + }, + protocols::worker_spec::WorkerConfigRequest, + workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, +}; + +/// Step 3: Create worker object(s) with merged configuration + metadata. +/// +/// This step: +/// 1. Merges discovered labels with config labels +/// 2. Determines the model ID from various sources +/// 3. Creates ModelCard with metadata +/// 4. Builds worker(s) - either single worker or multiple DP-aware workers +/// 5. Outputs unified `workers: Vec>` for downstream steps +pub struct CreateLocalWorkerStep; + +#[async_trait] +impl StepExecutor for CreateLocalWorkerStep { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + let config: Arc = context.get_or_err("worker_config")?; + let app_context: Arc = context.get_or_err("app_context")?; + let connection_mode: Arc = context.get_or_err("connection_mode")?; + let discovered_labels: Arc> = + context.get_or_err("discovered_labels")?; + + // Check if worker already exists + if app_context + .worker_registry + .get_by_url(&config.url) + .is_some() + { + return Err(WorkflowError::StepFailed { + step_id: StepId::new("create_worker"), + message: format!("Worker {} already exists", config.url), + }); + } + + // 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()); + } + if let Some(cost) = config.cost { + config_labels.insert("cost".to_string(), cost.to_string()); + } + + // Merge: discovered labels first, then config labels (config takes precedence) + let mut final_labels = discovered_labels.as_ref().clone(); + for (key, value) in &config_labels { + final_labels.insert(key.clone(), value.clone()); + } + + // Determine model_id: config > served_model_name > model_path > "unknown" + let model_id = config + .model_id + .clone() + .or_else(|| final_labels.get("served_model_name").cloned()) + .or_else(|| final_labels.get("model_path").cloned()) + .unwrap_or_else(|| "unknown".to_string()); + + if model_id != "unknown" { + debug!("Using model_id: {}", model_id); + } + + // Create ModelCard + let model_card = build_model_card(&model_id, &config, &final_labels); + + debug!( + "Creating worker {} with {} discovered + {} config = {} final labels", + config.url, + discovered_labels.len(), + config_labels.len(), + final_labels.len() + ); + + // Parse worker type + let worker_type = parse_worker_type(&config); + + // Get runtime type (for gRPC workers) + let runtime_type = determine_runtime_type(&connection_mode, context, &config); + + // Build circuit breaker config + let circuit_breaker_config = build_circuit_breaker_config(&app_context); + + // Build health config + let health_config = build_health_config(&app_context); + + // Normalize URL + let normalized_url = normalize_url(&config.url, &connection_mode); + + if normalized_url != config.url { + debug!( + "Normalized worker URL: {} -> {} ({:?})", + config.url, + normalized_url, + connection_mode.as_ref() + ); + } + + // Create workers - always output as Vec for unified downstream handling + let workers = if config.dp_aware { + create_dp_aware_workers( + context, + &normalized_url, + model_card, + worker_type, + &connection_mode, + runtime_type, + circuit_breaker_config, + health_config, + &config, + &final_labels, + )? + } else { + create_single_worker( + &normalized_url, + model_card, + worker_type, + &connection_mode, + runtime_type, + circuit_breaker_config, + health_config, + &config, + &final_labels, + ) + }; + + context.set("workers", workers); + context.set("labels", final_labels); + Ok(StepResult::Success) + } + + fn is_retryable(&self, _error: &WorkflowError) -> bool { + false + } +} + +fn build_model_card( + model_id: &str, + config: &WorkerConfigRequest, + labels: &HashMap, +) -> ModelCard { + let mut card = ModelCard::new(model_id); + + if let Some(ref tokenizer_path) = config.tokenizer_path { + card = card.with_tokenizer_path(tokenizer_path.clone()); + } + if let Some(ref reasoning_parser) = config.reasoning_parser { + card = card.with_reasoning_parser(reasoning_parser.clone()); + } + if let Some(ref tool_parser) = config.tool_parser { + card = card.with_tool_parser(tool_parser.clone()); + } + if let Some(ref chat_template) = config.chat_template { + card = card.with_chat_template(chat_template.clone()); + } + if let Some(model_type_str) = labels.get("model_type") { + card = card.with_hf_model_type(model_type_str.clone()); + } + if let Some(architectures_json) = labels.get("architectures") { + if let Ok(architectures) = serde_json::from_str::>(architectures_json) { + card = card.with_architectures(architectures); + } + } + + card +} + +fn parse_worker_type(config: &WorkerConfigRequest) -> WorkerType { + config + .worker_type + .as_ref() + .map(|t| match t.as_str() { + "prefill" => WorkerType::Prefill { + bootstrap_port: config.bootstrap_port, + }, + "decode" => WorkerType::Decode, + _ => WorkerType::Regular, + }) + .unwrap_or(WorkerType::Regular) +} + +fn determine_runtime_type( + connection_mode: &ConnectionMode, + context: &WorkflowContext, + config: &WorkerConfigRequest, +) -> RuntimeType { + if !matches!(connection_mode, ConnectionMode::Grpc { .. }) { + return RuntimeType::Sglang; + } + + if let Some(detected_runtime) = context.get::("detected_runtime_type") { + match detected_runtime.as_str() { + "vllm" => RuntimeType::Vllm, + _ => RuntimeType::Sglang, + } + } else if let Some(ref runtime) = config.runtime { + match runtime.as_str() { + "vllm" => RuntimeType::Vllm, + _ => RuntimeType::Sglang, + } + } else { + RuntimeType::Sglang + } +} + +fn build_circuit_breaker_config(app_context: &AppContext) -> CircuitBreakerConfig { + 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), + } +} + +fn build_health_config(app_context: &AppContext) -> HealthConfig { + 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, + } +} + +fn normalize_url(url: &str, connection_mode: &ConnectionMode) -> String { + if url.starts_with("http://") || url.starts_with("https://") || url.starts_with("grpc://") { + url.to_string() + } else { + match connection_mode { + ConnectionMode::Http => format!("http://{}", url), + ConnectionMode::Grpc { .. } => format!("grpc://{}", url), + } + } +} + +#[allow(clippy::too_many_arguments)] +fn create_dp_aware_workers( + context: &WorkflowContext, + normalized_url: &str, + model_card: ModelCard, + worker_type: WorkerType, + connection_mode: &ConnectionMode, + runtime_type: RuntimeType, + circuit_breaker_config: CircuitBreakerConfig, + health_config: HealthConfig, + config: &WorkerConfigRequest, + final_labels: &HashMap, +) -> Result>, WorkflowError> { + let dp_info: Arc = context.get_or_err("dp_info")?; + + debug!( + "Creating {} DP-aware workers for {} (dp_size: {})", + dp_info.dp_size, normalized_url, dp_info.dp_size + ); + + let mut workers = Vec::with_capacity(dp_info.dp_size); + for rank in 0..dp_info.dp_size { + let mut builder = + DPAwareWorkerBuilder::new(normalized_url.to_string(), rank, dp_info.dp_size) + .model(model_card.clone()) + .worker_type(worker_type.clone()) + .connection_mode(connection_mode.clone()) + .runtime_type(runtime_type.clone()) + .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 !final_labels.is_empty() { + builder = builder.labels(final_labels.clone()); + } + + let worker = Arc::new(builder.build()) as Arc; + worker.set_healthy(false); + workers.push(worker); + + debug!( + "Created DP-aware worker {}@{}/{} ({:?})", + normalized_url, rank, dp_info.dp_size, connection_mode + ); + } + + Ok(workers) +} + +#[allow(clippy::too_many_arguments)] +fn create_single_worker( + normalized_url: &str, + model_card: ModelCard, + worker_type: WorkerType, + connection_mode: &ConnectionMode, + runtime_type: RuntimeType, + circuit_breaker_config: CircuitBreakerConfig, + health_config: HealthConfig, + config: &WorkerConfigRequest, + final_labels: &HashMap, +) -> Vec> { + let mut builder = BasicWorkerBuilder::new(normalized_url.to_string()) + .model(model_card) + .worker_type(worker_type) + .connection_mode(connection_mode.clone()) + .runtime_type(runtime_type) + .circuit_breaker_config(circuit_breaker_config) + .health_config(health_config); + + 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()); + } + + let worker = Arc::new(builder.build()) as Arc; + worker.set_healthy(false); + + debug!( + "Created worker object for {} ({:?}) with {} labels", + normalized_url, + connection_mode, + final_labels.len() + ); + + vec![worker] +} diff --git a/sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs b/sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs new file mode 100644 index 000000000..697fb59d0 --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs @@ -0,0 +1,146 @@ +//! Connection mode detection step. + +use std::{sync::Arc, time::Duration}; + +use async_trait::async_trait; +use reqwest::Client; +use tracing::debug; + +use crate::{ + app_context::AppContext, + core::ConnectionMode, + protocols::worker_spec::WorkerConfigRequest, + routers::grpc::client::GrpcClient, + workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, +}; + +/// Strip protocol prefix from URL. +fn strip_protocol(url: &str) -> String { + url.trim_start_matches("http://") + .trim_start_matches("https://") + .trim_start_matches("grpc://") + .to_string() +} + +/// Try HTTP health check. +async fn try_http_health_check( + url: &str, + timeout_secs: u64, + client: &Client, +) -> Result<(), String> { + 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); + + client + .get(&health_url) + .timeout(Duration::from_secs(timeout_secs)) + .send() + .await + .and_then(reqwest::Response::error_for_status) + .map_err(|e| format!("Health check failed: {}", e))?; + + Ok(()) +} + +/// Perform gRPC health check with runtime type. +async fn do_grpc_health_check( + grpc_url: &str, + timeout_secs: u64, + runtime_type: &str, +) -> Result<(), String> { + let connect_future = GrpcClient::connect(grpc_url, runtime_type); + let client = tokio::time::timeout(Duration::from_secs(timeout_secs), connect_future) + .await + .map_err(|_| "gRPC connection timeout".to_string())? + .map_err(|e| format!("gRPC connection failed: {}", e))?; + + let health_future = client.health_check(); + tokio::time::timeout(Duration::from_secs(timeout_secs), health_future) + .await + .map_err(|_| "gRPC health check timeout".to_string())? + .map_err(|e| format!("gRPC health check failed: {}", e))?; + + Ok(()) +} + +/// Try gRPC health check (tries SGLang first, then vLLM if not specified). +async fn try_grpc_health_check( + url: &str, + timeout_secs: u64, + runtime_type: Option<&str>, +) -> Result<(), String> { + let grpc_url = if url.starts_with("grpc://") { + url.to_string() + } else { + format!("grpc://{}", strip_protocol(url)) + }; + + match runtime_type { + Some(runtime) => do_grpc_health_check(&grpc_url, timeout_secs, runtime).await, + None => { + // Try SGLang first, then vLLM as fallback + if let Ok(()) = do_grpc_health_check(&grpc_url, timeout_secs, "sglang").await { + return Ok(()); + } + do_grpc_health_check(&grpc_url, timeout_secs, "vllm") + .await + .map_err(|e| format!("gRPC failed (tried SGLang and vLLM): {}", e)) + } + } +} + +/// Step 1: Detect connection mode by probing HTTP and gRPC. +pub struct DetectConnectionModeStep; + +#[async_trait] +impl StepExecutor for DetectConnectionModeStep { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + let config: Arc = context.get_or_err("worker_config")?; + let app_context: Arc = context.get_or_err("app_context")?; + + debug!( + "Detecting connection mode for {} (timeout: {}s, max_attempts: {})", + config.url, config.health_check_timeout_secs, config.max_connection_attempts + ); + + // 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) + ); + + let connection_mode = match (http_result, grpc_result) { + (Ok(_), _) => { + debug!("{} detected as HTTP", config.url); + ConnectionMode::Http + } + (_, Ok(_)) => { + debug!("{} detected as gRPC", config.url); + ConnectionMode::Grpc { port: None } + } + (Err(http_err), Err(grpc_err)) => { + return Err(WorkflowError::StepFailed { + step_id: StepId::new("detect_connection_mode"), + message: format!( + "Both HTTP and gRPC health checks failed for {}: HTTP: {}, gRPC: {}", + config.url, http_err, grpc_err + ), + }); + } + }; + + context.set("connection_mode", connection_mode); + Ok(StepResult::Success) + } + + fn is_retryable(&self, _error: &WorkflowError) -> bool { + true + } +} diff --git a/sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs b/sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs new file mode 100644 index 000000000..33e1391cd --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs @@ -0,0 +1,79 @@ +//! Data Parallel (DP) information discovery step. + +use std::sync::Arc; + +use async_trait::async_trait; +use tracing::debug; + +use super::discover_metadata::get_server_info; +use crate::{ + protocols::worker_spec::WorkerConfigRequest, + workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, +}; + +/// DP (Data Parallel) information for a worker. +#[derive(Debug, Clone)] +pub struct DpInfo { + pub dp_size: usize, + pub model_id: String, +} + +/// Get DP info for a worker URL. +pub async fn get_dp_info(url: &str, api_key: Option<&str>) -> Result { + let info = get_server_info(url, api_key).await?; + + let dp_size = info + .dp_size + .ok_or_else(|| format!("No dp_size in response from {}", url))?; + + 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())) + }) + .unwrap_or_else(|| "unknown".to_string()); + + Ok(DpInfo { dp_size, model_id }) +} + +/// Step 2b: Discover DP (Data Parallel) information (only for DP-aware workers). +pub struct DiscoverDPInfoStep; + +#[async_trait] +impl StepExecutor for DiscoverDPInfoStep { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + let config: Arc = context.get_or_err("worker_config")?; + + if !config.dp_aware { + debug!( + "Worker {} is not DP-aware, skipping DP discovery", + config.url + ); + return Ok(StepResult::Success); + } + + debug!("Discovering DP info for {} (DP-aware)", config.url); + + let dp_info = get_dp_info(&config.url, config.api_key.as_deref()) + .await + .map_err(|e| WorkflowError::StepFailed { + step_id: StepId::new("discover_dp_info"), + message: format!("Failed to get DP info: {}", e), + })?; + + debug!( + "Discovered DP size {} for {} (model: {})", + dp_info.dp_size, config.url, dp_info.model_id + ); + + context.set("dp_info", dp_info); + Ok(StepResult::Success) + } + + fn is_retryable(&self, _error: &WorkflowError) -> bool { + true + } +} 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 new file mode 100644 index 000000000..62c4d47fa --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs @@ -0,0 +1,244 @@ +//! Metadata discovery step for local workers. + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use async_trait::async_trait; +use once_cell::sync::Lazy; +use reqwest::Client; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use tracing::{debug, warn}; + +use crate::{ + core::ConnectionMode, + protocols::worker_spec::WorkerConfigRequest, + routers::grpc::client::GrpcClient, + workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, +}; + +// HTTP client for metadata fetching +static HTTP_CLIENT: Lazy = Lazy::new(|| { + Client::builder() + .timeout(Duration::from_secs(10)) + .build() + .expect("Failed to create HTTP client") +}); + +/// Server information returned from /server_info endpoint. +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ServerInfo { + #[serde(alias = "model")] + pub model_id: Option, + pub model_path: Option, + pub served_model_name: Option, + pub dp_size: Option, + pub version: Option, + pub max_batch_size: Option, + pub max_total_tokens: Option, + pub max_prefill_tokens: Option, + pub max_running_requests: Option, + pub max_num_reqs: Option, +} + +/// Model information returned from /model_info endpoint. +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ModelInfo { + pub model_path: Option, + pub tokenizer_path: Option, + pub is_generation: Option, + pub model_type: Option, + pub architectures: Option>, +} + +/// Strip protocol prefix from URL. +fn strip_protocol(url: &str) -> String { + url.trim_start_matches("http://") + .trim_start_matches("https://") + .trim_start_matches("grpc://") + .to_string() +} + +/// Get server info from /server_info endpoint. +pub async fn get_server_info(url: &str, api_key: Option<&str>) -> Result { + let base_url = url.trim_end_matches('/'); + let server_info_url = format!("{}/server_info", base_url); + + let mut req = HTTP_CLIENT.get(&server_info_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 {}: {}", server_info_url, e))?; + + if !response.status().is_success() { + return Err(format!( + "Server returned status {} from {}", + response.status(), + server_info_url + )); + } + + let json = response + .json::() + .await + .map_err(|e| format!("Failed to parse response from {}: {}", server_info_url, e))?; + + serde_json::from_value(json).map_err(|e| format!("Failed to parse server info: {}", e)) +} + +/// Get model info from /model_info endpoint. +pub async fn get_model_info(url: &str, api_key: Option<&str>) -> Result { + let base_url = url.trim_end_matches('/'); + let model_info_url = format!("{}/model_info", base_url); + + let mut req = HTTP_CLIENT.get(&model_info_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 {}: {}", model_info_url, e))?; + + if !response.status().is_success() { + return Err(format!( + "Server returned status {} from {}", + response.status(), + model_info_url + )); + } + + response + .json::() + .await + .map_err(|e| format!("Failed to parse response from {}: {}", model_info_url, e)) +} + +/// Fetch gRPC metadata (returns labels and detected runtime type). +async fn fetch_grpc_metadata( + url: &str, + runtime_type: Option<&str>, +) -> Result<(HashMap, String), String> { + let grpc_url = if url.starts_with("grpc://") { + url.to_string() + } else { + format!("grpc://{}", strip_protocol(url)) + }; + + async fn do_fetch( + grpc_url: &str, + runtime_type: &str, + ) -> Result, String> { + let client = GrpcClient::connect(grpc_url, runtime_type) + .await + .map_err(|e| format!("Failed to connect to gRPC: {}", e))?; + + let model_info = client + .get_model_info() + .await + .map_err(|e| format!("Failed to fetch gRPC metadata: {}", e))?; + + Ok(model_info.to_labels()) + } + + match runtime_type { + Some(runtime) => { + let labels = do_fetch(&grpc_url, runtime).await?; + Ok((labels, runtime.to_string())) + } + None => { + // Try SGLang first, then vLLM as fallback + if let Ok(labels) = do_fetch(&grpc_url, "sglang").await { + return Ok((labels, "sglang".to_string())); + } + let labels = do_fetch(&grpc_url, "vllm") + .await + .map_err(|e| format!("gRPC metadata failed (tried SGLang and vLLM): {}", e))?; + Ok((labels, "vllm".to_string())) + } + } +} + +/// Step 2a: Discover metadata from worker. +pub struct DiscoverMetadataStep; + +#[async_trait] +impl StepExecutor for DiscoverMetadataStep { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + let config: Arc = context.get_or_err("worker_config")?; + let connection_mode: Arc = context.get_or_err("connection_mode")?; + + debug!( + "Discovering metadata for {} ({:?})", + config.url, *connection_mode + ); + + let (discovered_labels, detected_runtime) = match connection_mode.as_ref() { + ConnectionMode::Http => { + let mut labels = HashMap::new(); + + // Fetch from /server_info for server-related metadata + if let Ok(server_info) = + get_server_info(&config.url, config.api_key.as_deref()).await + { + if let Some(model_path) = server_info.model_path.filter(|s| !s.is_empty()) { + labels.insert("model_path".to_string(), model_path); + } + if let Some(served_model_name) = + server_info.served_model_name.filter(|s| !s.is_empty()) + { + labels.insert("served_model_name".to_string(), served_model_name); + } + } + + // Fetch from /model_info for model-related metadata + if let Ok(model_info) = get_model_info(&config.url, config.api_key.as_deref()).await + { + if let Some(model_type) = model_info.model_type.filter(|s| !s.is_empty()) { + labels.insert("model_type".to_string(), model_type); + } + if let Some(architectures) = model_info.architectures.filter(|a| !a.is_empty()) + { + if let Ok(json_str) = serde_json::to_string(&architectures) { + labels.insert("architectures".to_string(), json_str); + } + } + } + + Ok((labels, None)) + } + ConnectionMode::Grpc { .. } => { + let runtime_type = config.runtime.as_deref(); + fetch_grpc_metadata(&config.url, runtime_type) + .await + .map(|(labels, runtime)| (labels, Some(runtime))) + } + } + .unwrap_or_else(|e| { + warn!("Failed to fetch metadata for {}: {}", config.url, e); + (HashMap::new(), None) + }); + + debug!( + "Discovered {} metadata labels for {}", + discovered_labels.len(), + config.url + ); + + context.set("discovered_labels", discovered_labels); + if let Some(runtime) = detected_runtime { + debug!("Detected runtime type: {}", runtime); + context.set("detected_runtime_type", runtime); + } + + Ok(StepResult::Success) + } + + fn is_retryable(&self, _error: &WorkflowError) -> bool { + true + } +} diff --git a/sgl-model-gateway/src/core/steps/worker/local/mod.rs b/sgl-model-gateway/src/core/steps/worker/local/mod.rs new file mode 100644 index 000000000..3a2890219 --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/local/mod.rs @@ -0,0 +1,195 @@ +mod create_worker; +mod detect_connection; +mod discover_dp; +mod discover_metadata; +mod removal; + +use std::{sync::Arc, time::Duration}; + +pub use create_worker::CreateLocalWorkerStep; +pub use detect_connection::DetectConnectionModeStep; +pub use discover_dp::{get_dp_info, DiscoverDPInfoStep, DpInfo}; +pub use discover_metadata::DiscoverMetadataStep; +pub use removal::{ + FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep, RemoveFromWorkerRegistryStep, + UpdateRemainingPoliciesStep, WorkerRemovalRequest, +}; + +use super::shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep}; +use crate::{ + config::RouterConfig, + workflow::{BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition}, +}; + +pub fn create_local_worker_workflow(router_config: &RouterConfig) -> WorkflowDefinition { + let detect_timeout = Duration::from_secs(router_config.worker_startup_timeout_secs); + + // 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 { + (5 + ((effective_timeout - 10.0) / 5.0).ceil() as u32).max(3) + } else { + 3 + }; + + WorkflowDefinition::new("local_worker_registration", "Local Worker Registration") + // Step 1: Detect connection mode (HTTP vs gRPC) + .add_step( + StepDefinition::new( + "detect_connection_mode", + "Detect Connection Mode", + Arc::new(DetectConnectionModeStep), + ) + .with_retry(RetryPolicy { + max_attempts, + backoff: BackoffStrategy::Linear { + increment: Duration::from_secs(1), + max: Duration::from_secs(5), + }, + }) + .with_timeout(detect_timeout) + .with_failure_action(FailureAction::FailWorkflow), + ) + // Step 2a: Discover metadata (parallel with DP discovery) + .add_step( + StepDefinition::new( + "discover_metadata", + "Discover Metadata", + Arc::new(DiscoverMetadataStep), + ) + .with_retry(RetryPolicy { + max_attempts: 3, + backoff: BackoffStrategy::Fixed(Duration::from_secs(1)), + }) + .with_timeout(Duration::from_secs(10)) + .with_failure_action(FailureAction::ContinueNextStep) + .depends_on(&["detect_connection_mode"]), + ) + // Step 2b: Discover DP info (after metadata to avoid concurrent /server_info calls) + .add_step( + StepDefinition::new( + "discover_dp_info", + "Discover DP Info", + Arc::new(DiscoverDPInfoStep), + ) + .with_retry(RetryPolicy { + max_attempts: 3, + backoff: BackoffStrategy::Fixed(Duration::from_secs(1)), + }) + .with_timeout(Duration::from_secs(10)) + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["discover_metadata"]), + ) + // Step 3: Create worker(s) + .add_step( + StepDefinition::new( + "create_worker", + "Create Worker", + Arc::new(CreateLocalWorkerStep), + ) + .with_timeout(Duration::from_secs(5)) + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["discover_dp_info"]), + ) + // Step 4: Register workers (shared step) + .add_step( + StepDefinition::new( + "register_workers", + "Register Workers", + Arc::new(RegisterWorkersStep), + ) + .with_timeout(Duration::from_secs(5)) + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["create_worker"]), + ) + // Step 5a: Update policies (parallel with activation) + .add_step( + StepDefinition::new( + "update_policies", + "Update Policies", + Arc::new(UpdatePoliciesStep), + ) + .with_timeout(Duration::from_secs(5)) + .with_failure_action(FailureAction::ContinueNextStep) + .depends_on(&["register_workers"]), + ) + // Step 5b: Activate workers (parallel with policy update) + .add_step( + StepDefinition::new( + "activate_workers", + "Activate Workers", + Arc::new(ActivateWorkersStep), + ) + .with_timeout(Duration::from_secs(5)) + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["register_workers"]), + ) +} + +/// Create a worker removal workflow definition. +/// +/// DAG structure: +/// ```text +/// find_workers_to_remove +/// │ +/// remove_from_policy_registry +/// │ +/// remove_from_worker_registry +/// │ +/// update_remaining_policies +/// ``` +pub fn create_worker_removal_workflow() -> WorkflowDefinition { + WorkflowDefinition::new("worker_removal", "Remove worker from router") + .add_step( + StepDefinition::new( + "find_workers_to_remove", + "Find workers to remove", + Arc::new(FindWorkersToRemoveStep), + ) + .with_timeout(Duration::from_secs(10)) + .with_retry(RetryPolicy { + max_attempts: 1, + backoff: BackoffStrategy::Fixed(Duration::from_secs(0)), + }), + ) + .add_step( + StepDefinition::new( + "remove_from_policy_registry", + "Remove workers from policy registry", + Arc::new(RemoveFromPolicyRegistryStep), + ) + .with_timeout(Duration::from_secs(10)) + .with_retry(RetryPolicy { + max_attempts: 1, + backoff: BackoffStrategy::Fixed(Duration::from_secs(0)), + }) + .depends_on(&["find_workers_to_remove"]), + ) + .add_step( + StepDefinition::new( + "remove_from_worker_registry", + "Remove workers from worker registry", + Arc::new(RemoveFromWorkerRegistryStep), + ) + .with_timeout(Duration::from_secs(10)) + .with_retry(RetryPolicy { + max_attempts: 1, + backoff: BackoffStrategy::Fixed(Duration::from_secs(0)), + }) + .depends_on(&["remove_from_policy_registry"]), + ) + .add_step( + StepDefinition::new( + "update_remaining_policies", + "Update cache-aware policies for remaining workers", + Arc::new(UpdateRemainingPoliciesStep), + ) + .with_timeout(Duration::from_secs(10)) + .with_retry(RetryPolicy { + max_attempts: 1, + backoff: BackoffStrategy::Fixed(Duration::from_secs(0)), + }) + .depends_on(&["remove_from_worker_registry"]), + ) +} diff --git a/sgl-model-gateway/src/core/steps/worker_removal.rs b/sgl-model-gateway/src/core/steps/worker/local/removal.rs similarity index 71% rename from sgl-model-gateway/src/core/steps/worker_removal.rs rename to sgl-model-gateway/src/core/steps/worker/local/removal.rs index dad96fa90..d94e243ee 100644 --- a/sgl-model-gateway/src/core/steps/worker_removal.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/removal.rs @@ -1,18 +1,24 @@ -use std::{collections::HashSet, sync::Arc, time::Duration}; +//! Worker removal step implementations. + +use std::{collections::HashSet, sync::Arc}; use async_trait::async_trait; use tracing::{debug, info}; -use crate::{app_context::AppContext, core::Worker, workflow::*}; +use crate::{ + app_context::AppContext, + core::Worker, + workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, +}; -/// Request structure for worker removal +/// Request structure for worker removal. #[derive(Debug, Clone)] pub struct WorkerRemovalRequest { pub url: String, pub dp_aware: bool, } -/// Step 1: Find workers to remove based on URL +/// Step 1: Find workers to remove based on URL. pub struct FindWorkersToRemoveStep; #[async_trait] @@ -81,11 +87,11 @@ impl StepExecutor for FindWorkersToRemoveStep { } fn is_retryable(&self, _error: &WorkflowError) -> bool { - false // Worker not found is not retryable + false } } -/// Step 2: Remove workers from policy registry +/// Step 2: Remove workers from policy registry. pub struct RemoveFromPolicyRegistryStep; #[async_trait] @@ -122,11 +128,11 @@ impl StepExecutor for RemoveFromPolicyRegistryStep { } fn is_retryable(&self, _error: &WorkflowError) -> bool { - false // Policy removal is not retryable + false } } -/// Step 3: Remove workers from worker registry +/// Step 3: Remove workers from worker registry. pub struct RemoveFromWorkerRegistryStep; #[async_trait] @@ -167,11 +173,11 @@ impl StepExecutor for RemoveFromWorkerRegistryStep { } fn is_retryable(&self, _error: &WorkflowError) -> bool { - false // Worker removal is not retryable + false } } -/// Step 4: Update cache-aware policies for remaining workers +/// Step 4: Update cache-aware policies for remaining workers. pub struct UpdateRemainingPoliciesStep; #[async_trait] @@ -219,62 +225,6 @@ impl StepExecutor for UpdateRemainingPoliciesStep { } fn is_retryable(&self, _error: &WorkflowError) -> bool { - false // Policy update is not retryable + false } } - -/// Create a worker removal workflow definition -pub fn create_worker_removal_workflow() -> WorkflowDefinition { - WorkflowDefinition::new("worker_removal", "Remove worker from router") - .add_step( - StepDefinition::new( - "find_workers_to_remove", - "Find workers to remove", - Arc::new(FindWorkersToRemoveStep), - ) - .with_timeout(Duration::from_secs(10)) - .with_retry(RetryPolicy { - max_attempts: 1, - backoff: BackoffStrategy::Fixed(Duration::from_secs(0)), - }), - ) - .add_step( - StepDefinition::new( - "remove_from_policy_registry", - "Remove workers from policy registry", - Arc::new(RemoveFromPolicyRegistryStep), - ) - .with_timeout(Duration::from_secs(10)) - .with_retry(RetryPolicy { - max_attempts: 1, - backoff: BackoffStrategy::Fixed(Duration::from_secs(0)), - }) - .depends_on(&["find_workers_to_remove"]), - ) - .add_step( - StepDefinition::new( - "remove_from_worker_registry", - "Remove workers from worker registry", - Arc::new(RemoveFromWorkerRegistryStep), - ) - .with_timeout(Duration::from_secs(10)) - .with_retry(RetryPolicy { - max_attempts: 1, - backoff: BackoffStrategy::Fixed(Duration::from_secs(0)), - }) - .depends_on(&["remove_from_policy_registry"]), - ) - .add_step( - StepDefinition::new( - "update_remaining_policies", - "Update cache-aware policies for remaining workers", - Arc::new(UpdateRemainingPoliciesStep), - ) - .with_timeout(Duration::from_secs(10)) - .with_retry(RetryPolicy { - max_attempts: 1, - backoff: BackoffStrategy::Fixed(Duration::from_secs(0)), - }) - .depends_on(&["remove_from_worker_registry"]), - ) -} diff --git a/sgl-model-gateway/src/core/steps/worker/mod.rs b/sgl-model-gateway/src/core/steps/worker/mod.rs new file mode 100644 index 000000000..dd8b1fc8a --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/mod.rs @@ -0,0 +1,17 @@ +pub mod external; +pub mod local; +pub mod shared; + +pub use external::{ + create_external_worker_workflow as create_external_worker_registration_workflow, + create_external_worker_workflow, group_models_into_cards, infer_model_type_from_id, + CreateExternalWorkersStep, DiscoverModelsStep, ModelInfo, ModelsResponse, +}; +pub use local::{ + create_local_worker_workflow as create_worker_registration_workflow, + create_local_worker_workflow, create_worker_removal_workflow, CreateLocalWorkerStep, + DetectConnectionModeStep, DiscoverDPInfoStep, DiscoverMetadataStep, DpInfo, + FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep, RemoveFromWorkerRegistryStep, + UpdateRemainingPoliciesStep, WorkerRemovalRequest, +}; +pub use shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep, WorkerList}; diff --git a/sgl-model-gateway/src/core/steps/worker/shared/activate.rs b/sgl-model-gateway/src/core/steps/worker/shared/activate.rs new file mode 100644 index 000000000..eaee82962 --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/shared/activate.rs @@ -0,0 +1,35 @@ +//! Unified worker activation step. + +use std::sync::Arc; + +use async_trait::async_trait; +use tracing::info; + +use crate::{ + core::Worker, + workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowResult}, +}; + +/// Unified step to activate workers by marking them as healthy. +/// +/// This is the final step in any worker registration workflow. +pub struct ActivateWorkersStep; + +#[async_trait] +impl StepExecutor for ActivateWorkersStep { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + let workers: Arc>> = context.get_or_err("workers")?; + + for worker in workers.iter() { + worker.set_healthy(true); + } + + info!("Activated {} worker(s) (marked as healthy)", workers.len()); + + Ok(StepResult::Success) + } + + fn is_retryable(&self, _error: &crate::workflow::WorkflowError) -> bool { + false + } +} diff --git a/sgl-model-gateway/src/core/steps/worker/shared/mod.rs b/sgl-model-gateway/src/core/steps/worker/shared/mod.rs new file mode 100644 index 000000000..f60877c6f --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/shared/mod.rs @@ -0,0 +1,21 @@ +//! Shared worker registration steps used by both local and external workflows. +//! +//! These steps are designed to work with any worker type and can be composed +//! into different workflows using the DAG-based workflow engine. + +mod activate; +mod register; +mod update_policies; + +use std::sync::Arc; + +pub use activate::ActivateWorkersStep; +pub use register::RegisterWorkersStep; +pub use update_policies::UpdatePoliciesStep; + +use crate::core::Worker; + +/// Type alias for a collection of workers in workflow context. +/// Both local (single/DP-aware) and external (multi-model) workflows +/// use this unified type for consistency. +pub type WorkerList = Vec>; diff --git a/sgl-model-gateway/src/core/steps/worker/shared/register.rs b/sgl-model-gateway/src/core/steps/worker/shared/register.rs new file mode 100644 index 000000000..b63962647 --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/shared/register.rs @@ -0,0 +1,46 @@ +//! Unified worker registration step. + +use std::sync::Arc; + +use async_trait::async_trait; +use tracing::debug; + +use crate::{ + app_context::AppContext, + core::Worker, + workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowResult}, +}; + +/// Unified step to register workers in the registry. +/// +/// Works with both single workers and batches. Always expects `workers` key +/// in context containing `Vec>`. +pub struct RegisterWorkersStep; + +#[async_trait] +impl StepExecutor for RegisterWorkersStep { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + let app_context: Arc = context.get_or_err("app_context")?; + let workers: Arc>> = context.get_or_err("workers")?; + + let mut worker_ids = Vec::with_capacity(workers.len()); + + for worker in workers.iter() { + let worker_id = app_context.worker_registry.register(Arc::clone(worker)); + debug!( + "Registered worker {} (model: {}) with ID {:?}", + worker.url(), + worker.model_id(), + worker_id + ); + worker_ids.push(worker_id); + } + + context.set("worker_ids", worker_ids); + Ok(StepResult::Success) + } + + fn is_retryable(&self, _error: &crate::workflow::WorkflowError) -> bool { + false + } +} 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 new file mode 100644 index 000000000..282bd2972 --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs @@ -0,0 +1,78 @@ +//! Unified policy update step. + +use std::{collections::HashMap, sync::Arc}; + +use async_trait::async_trait; +use tracing::debug; + +use crate::{ + app_context::AppContext, + core::Worker, + workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowResult}, +}; + +/// Unified step to update policy registry for registered workers. +/// +/// Handles both local workers (same model, possibly DP-aware) and +/// external workers (different models per worker). +pub struct UpdatePoliciesStep; + +#[async_trait] +impl StepExecutor for UpdatePoliciesStep { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + let app_context: Arc = context.get_or_err("app_context")?; + let workers: Arc>> = context.get_or_err("workers")?; + let labels: Arc> = context.get_or_err("labels")?; + + let policy_hint = labels.get("policy").map(|s| s.as_str()); + + // Track unique model IDs we've updated policies for + let mut updated_models = Vec::new(); + + for worker in workers.iter() { + let model_id = worker.model_id().to_string(); + + // Notify policy registry + app_context + .policy_registry + .on_worker_added(&model_id, policy_hint); + + // Initialize cache-aware policy if configured + let all_workers = app_context.worker_registry.get_by_model_fast(&model_id); + if let Some(policy) = app_context.policy_registry.get_policy(&model_id) { + if policy.name() == "cache_aware" { + app_context + .policy_registry + .init_cache_aware_policy(&model_id, &all_workers); + } + } + + if !updated_models.contains(&model_id) { + updated_models.push(model_id); + } + } + + // Initialize bucket policies for prefill workers (local workers only) + let prefill_workers = app_context.worker_registry.get_prefill_workers(); + if !prefill_workers.is_empty() { + let policy = app_context.policy_registry.get_prefill_policy(); + if policy.name() == "bucket" { + app_context + .policy_registry + .init_pd_bucket_policies(&prefill_workers); + } + } + + debug!( + "Updated policies for {} workers across {} models", + workers.len(), + updated_models.len() + ); + + Ok(StepResult::Success) + } + + fn is_retryable(&self, _error: &crate::workflow::WorkflowError) -> bool { + false + } +} diff --git a/sgl-model-gateway/src/core/steps/worker_registration.rs b/sgl-model-gateway/src/core/steps/worker_registration.rs deleted file mode 100644 index c936615a8..000000000 --- a/sgl-model-gateway/src/core/steps/worker_registration.rs +++ /dev/null @@ -1,943 +0,0 @@ -use std::{collections::HashMap, sync::Arc, time::Duration}; - -use async_trait::async_trait; -use once_cell::sync::Lazy; -use reqwest::Client; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use tracing::{debug, info, warn}; - -use crate::{ - app_context::AppContext, - core::{ - BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, DPAwareWorkerBuilder, - HealthConfig, ModelCard, RuntimeType, Worker, WorkerType, - }, - protocols::worker_spec::WorkerConfigRequest, - routers::grpc::client::GrpcClient, - workflow::*, -}; - -// HTTP client for metadata fetching -static HTTP_CLIENT: Lazy = Lazy::new(|| { - Client::builder() - .timeout(Duration::from_secs(10)) - .build() - .expect("Failed to create HTTP client") -}); - -/// Server information returned from /server_info endpoint -#[derive(Debug, Clone, Deserialize, Serialize)] -struct ServerInfo { - #[serde(alias = "model")] - model_id: Option, - model_path: Option, - served_model_name: Option, - dp_size: Option, - version: Option, - max_batch_size: Option, - max_total_tokens: Option, - max_prefill_tokens: Option, - max_running_requests: Option, - max_num_reqs: Option, -} - -/// Model information returned from /model_info endpoint -#[derive(Debug, Clone, Deserialize, Serialize)] -struct ModelInfo { - model_path: Option, - tokenizer_path: Option, - is_generation: Option, - /// HuggingFace model type string (e.g., "llama", "qwen2", "gpt_oss") - model_type: Option, - /// Model architectures from HuggingFace config (e.g., ["LlamaForCausalLM"]) - architectures: Option>, -} - -#[derive(Debug, Clone)] -pub struct DpInfo { - pub dp_size: usize, - pub model_id: String, -} - -/// 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)) -} - -/// Get server info from /server_info endpoint -async fn get_server_info(url: &str, api_key: Option<&str>) -> Result { - let base_url = url.trim_end_matches('/'); - let server_info_url = format!("{}/server_info", base_url); - - let mut req = HTTP_CLIENT.get(&server_info_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 {}: {}", server_info_url, e))?; - - if !response.status().is_success() { - return Err(format!( - "Server returned status {} from {}", - response.status(), - server_info_url - )); - } - - let json = response - .json::() - .await - .map_err(|e| format!("Failed to parse response from {}: {}", server_info_url, e))?; - - parse_server_info(json) -} - -/// Get model info from /model_info endpoint -async fn get_model_info(url: &str, api_key: Option<&str>) -> Result { - let base_url = url.trim_end_matches('/'); - let model_info_url = format!("{}/model_info", base_url); - - let mut req = HTTP_CLIENT.get(&model_info_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 {}: {}", model_info_url, e))?; - - if !response.status().is_success() { - return Err(format!( - "Server returned status {} from {}", - response.status(), - model_info_url - )); - } - - response - .json::() - .await - .map_err(|e| format!("Failed to parse response from {}: {}", model_info_url, e)) -} - -/// Get DP info for a worker URL -async fn get_dp_info(url: &str, api_key: Option<&str>) -> Result { - let info = get_server_info(url, api_key).await?; - - let dp_size = info - .dp_size - .ok_or_else(|| format!("No dp_size in response from {}", url))?; - - 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())) - }) - .unwrap_or_else(|| "unknown".to_string()); - - Ok(DpInfo { dp_size, model_id }) -} - -/// Strip protocol prefix from URL -fn strip_protocol(url: &str) -> String { - url.trim_start_matches("http://") - .trim_start_matches("https://") - .trim_start_matches("grpc://") - .to_string() -} - -/// Try HTTP health check -async fn try_http_health_check( - url: &str, - timeout_secs: u64, - client: &Client, -) -> Result<(), String> { - 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); - - client - .get(&health_url) - .timeout(Duration::from_secs(timeout_secs)) - .send() - .await - .and_then(reqwest::Response::error_for_status) - .map_err(|e| format!("Health check failed: {}", e))?; - - Ok(()) -} - -/// Perform gRPC health check with runtime type -async fn do_grpc_health_check( - grpc_url: &str, - timeout_secs: u64, - runtime_type: &str, -) -> Result<(), String> { - let connect_future = GrpcClient::connect(grpc_url, runtime_type); - let client = tokio::time::timeout(Duration::from_secs(timeout_secs), connect_future) - .await - .map_err(|_| "gRPC connection timeout".to_string())? - .map_err(|e| format!("gRPC connection failed: {}", e))?; - - let health_future = client.health_check(); - tokio::time::timeout(Duration::from_secs(timeout_secs), health_future) - .await - .map_err(|_| "gRPC health check timeout".to_string())? - .map_err(|e| format!("gRPC health check failed: {}", e))?; - - Ok(()) -} - -/// Try gRPC health check (tries SGLang first, then vLLM if not specified) -async fn try_grpc_health_check( - url: &str, - timeout_secs: u64, - runtime_type: Option<&str>, -) -> Result<(), String> { - let grpc_url = if url.starts_with("grpc://") { - url.to_string() - } else { - format!("grpc://{}", strip_protocol(url)) - }; - - match runtime_type { - Some(runtime) => do_grpc_health_check(&grpc_url, timeout_secs, runtime).await, - None => { - // Try SGLang first, then vLLM as fallback - if let Ok(()) = do_grpc_health_check(&grpc_url, timeout_secs, "sglang").await { - return Ok(()); - } - do_grpc_health_check(&grpc_url, timeout_secs, "vllm") - .await - .map_err(|e| format!("gRPC failed (tried SGLang and vLLM): {}", e)) - } - } -} - -/// Fetch metadata from gRPC server -async fn do_fetch_grpc_metadata( - grpc_url: &str, - runtime_type: &str, -) -> Result, String> { - let client = GrpcClient::connect(grpc_url, runtime_type) - .await - .map_err(|e| format!("Failed to connect to gRPC: {}", e))?; - - let model_info = client - .get_model_info() - .await - .map_err(|e| format!("Failed to fetch gRPC metadata: {}", e))?; - - Ok(model_info.to_labels()) -} - -/// Fetch gRPC metadata (returns labels and detected runtime type) -async fn fetch_grpc_metadata( - url: &str, - runtime_type: Option<&str>, -) -> Result<(HashMap, String), String> { - let grpc_url = if url.starts_with("grpc://") { - url.to_string() - } else { - format!("grpc://{}", strip_protocol(url)) - }; - - match runtime_type { - Some(runtime) => { - let labels = do_fetch_grpc_metadata(&grpc_url, runtime).await?; - Ok((labels, runtime.to_string())) - } - None => { - // 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())); - } - let labels = do_fetch_grpc_metadata(&grpc_url, "vllm") - .await - .map_err(|e| format!("gRPC metadata failed (tried SGLang and vLLM): {}", e))?; - Ok((labels, "vllm".to_string())) - } - } -} - -/// Step 1: Detect connection mode by probing HTTP and gRPC -pub struct DetectConnectionModeStep; - -#[async_trait] -impl StepExecutor for DetectConnectionModeStep { - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context.get_or_err("worker_config")?; - let app_context: Arc = context.get_or_err("app_context")?; - - debug!( - "Detecting connection mode for {} (timeout: {}s, max_attempts: {})", - config.url, config.health_check_timeout_secs, config.max_connection_attempts - ); - - // 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) - ); - - let connection_mode = match (http_result, grpc_result) { - (Ok(_), _) => { - debug!("{} detected as HTTP", config.url); - ConnectionMode::Http - } - (_, Ok(_)) => { - debug!("{} detected as gRPC", config.url); - ConnectionMode::Grpc { port: None } - } - (Err(http_err), Err(grpc_err)) => { - return Err(WorkflowError::StepFailed { - step_id: StepId::new("detect_connection_mode"), - message: format!( - "Both HTTP and gRPC health checks failed for {}: HTTP: {}, gRPC: {}", - config.url, http_err, grpc_err - ), - }); - } - }; - - context.set("connection_mode", connection_mode); - Ok(StepResult::Success) - } - - fn is_retryable(&self, _error: &WorkflowError) -> bool { - true - } -} - -/// Step 2: Discover metadata from worker -pub struct DiscoverMetadataStep; - -#[async_trait] -impl StepExecutor for DiscoverMetadataStep { - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context.get_or_err("worker_config")?; - let connection_mode: Arc = context.get_or_err("connection_mode")?; - - debug!( - "Discovering metadata for {} ({:?})", - config.url, *connection_mode - ); - - let (discovered_labels, detected_runtime) = match connection_mode.as_ref() { - ConnectionMode::Http => { - let mut labels = HashMap::new(); - - // Fetch from /server_info for server-related metadata - if let Ok(server_info) = - get_server_info(&config.url, config.api_key.as_deref()).await - { - if let Some(model_path) = server_info.model_path.filter(|s| !s.is_empty()) { - labels.insert("model_path".to_string(), model_path); - } - if let Some(served_model_name) = - server_info.served_model_name.filter(|s| !s.is_empty()) - { - labels.insert("served_model_name".to_string(), served_model_name); - } - } - - // Fetch from /model_info for model-related metadata (model_type, architectures) - if let Ok(model_info) = get_model_info(&config.url, config.api_key.as_deref()).await - { - if let Some(model_type) = model_info.model_type.filter(|s| !s.is_empty()) { - labels.insert("model_type".to_string(), model_type); - } - if let Some(architectures) = model_info.architectures.filter(|a| !a.is_empty()) - { - if let Ok(json_str) = serde_json::to_string(&architectures) { - labels.insert("architectures".to_string(), json_str); - } - } - } - - Ok((labels, None)) - } - ConnectionMode::Grpc { .. } => { - let runtime_type = config.runtime.as_deref(); - fetch_grpc_metadata(&config.url, runtime_type) - .await - .map(|(labels, runtime)| (labels, Some(runtime))) - } - } - .unwrap_or_else(|e| { - warn!("Failed to fetch metadata for {}: {}", config.url, e); - (HashMap::new(), None) - }); - - debug!( - "Discovered {} metadata labels for {}", - discovered_labels.len(), - config.url - ); - - context.set("discovered_labels", discovered_labels); - if let Some(runtime) = detected_runtime { - debug!("Detected runtime type: {}", runtime); - context.set("detected_runtime_type", runtime); - } - - Ok(StepResult::Success) - } - - fn is_retryable(&self, _error: &WorkflowError) -> bool { - true - } -} - -/// Step 2.5: Discover DP (Data Parallel) information (only for DP-aware workers) -pub struct DiscoverDPInfoStep; - -#[async_trait] -impl StepExecutor for DiscoverDPInfoStep { - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context.get_or_err("worker_config")?; - - if !config.dp_aware { - debug!( - "Worker {} is not DP-aware, skipping DP discovery", - config.url - ); - return Ok(StepResult::Success); - } - - debug!("Discovering DP info for {} (DP-aware)", config.url); - - let dp_info = get_dp_info(&config.url, config.api_key.as_deref()) - .await - .map_err(|e| WorkflowError::StepFailed { - step_id: StepId::new("discover_dp_info"), - message: format!("Failed to get DP info: {}", e), - })?; - - debug!( - "Discovered DP size {} for {} (model: {})", - dp_info.dp_size, config.url, dp_info.model_id - ); - - context.set("dp_info", dp_info); - Ok(StepResult::Success) - } - - fn is_retryable(&self, _error: &WorkflowError) -> bool { - true - } -} - -/// Step 3: Create worker object with merged configuration + metadata -pub struct CreateWorkerStep; - -#[async_trait] -impl StepExecutor for CreateWorkerStep { - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context.get_or_err("worker_config")?; - let app_context: Arc = context.get_or_err("app_context")?; - let connection_mode: Arc = context.get_or_err("connection_mode")?; - let discovered_labels: Arc> = - context.get_or_err("discovered_labels")?; - - // Check if worker already exists - if app_context - .worker_registry - .get_by_url(&config.url) - .is_some() - { - return Err(WorkflowError::StepFailed { - step_id: StepId::new("create_worker"), - message: format!("Worker {} already exists", config.url), - }); - } - - // 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()); - } - if let Some(cost) = config.cost { - config_labels.insert("cost".to_string(), cost.to_string()); - } - - // Merge: discovered labels first, then config labels (config takes precedence) - let mut final_labels = discovered_labels.as_ref().clone(); - for (key, value) in &config_labels { - final_labels.insert(key.clone(), value.clone()); - } - - // Determine model_id: config > served_model_name > model_path > "unknown" - let model_id = config - .model_id - .clone() - .or_else(|| final_labels.get("served_model_name").cloned()) - .or_else(|| final_labels.get("model_path").cloned()) - .unwrap_or_else(|| "unknown".to_string()); - - if model_id != "unknown" { - debug!("Using model_id: {}", model_id); - } - - // Create ModelCard - let model_card = { - let mut card = ModelCard::new(&model_id); - if let Some(ref tokenizer_path) = config.tokenizer_path { - card = card.with_tokenizer_path(tokenizer_path.clone()); - } - if let Some(ref reasoning_parser) = config.reasoning_parser { - card = card.with_reasoning_parser(reasoning_parser.clone()); - } - if let Some(ref tool_parser) = config.tool_parser { - card = card.with_tool_parser(tool_parser.clone()); - } - if let Some(ref chat_template) = config.chat_template { - card = card.with_chat_template(chat_template.clone()); - } - // Set HuggingFace model type from discovered labels - if let Some(model_type_str) = final_labels.get("model_type") { - card = card.with_hf_model_type(model_type_str.clone()); - } - // Set architectures from discovered labels (JSON array string) - if let Some(architectures_json) = final_labels.get("architectures") { - if let Ok(architectures) = serde_json::from_str::>(architectures_json) { - card = card.with_architectures(architectures); - } - } - card - }; - - debug!( - "Creating worker {} with {} discovered + {} config = {} final labels", - config.url, - discovered_labels.len(), - config_labels.len(), - final_labels.len() - ); - - // Parse worker type - let worker_type = config - .worker_type - .as_ref() - .map(|t| match t.as_str() { - "prefill" => WorkerType::Prefill { - bootstrap_port: config.bootstrap_port, - }, - "decode" => WorkerType::Decode, - _ => WorkerType::Regular, - }) - .unwrap_or(WorkerType::Regular); - - // Get runtime type (for gRPC workers) - let runtime_type = if matches!(connection_mode.as_ref(), ConnectionMode::Grpc { .. }) { - if let Some(detected_runtime) = context.get::("detected_runtime_type") { - match detected_runtime.as_str() { - "vllm" => RuntimeType::Vllm, - _ => RuntimeType::Sglang, - } - } else if let Some(ref runtime) = config.runtime { - match runtime.as_str() { - "vllm" => RuntimeType::Vllm, - _ => RuntimeType::Sglang, - } - } else { - RuntimeType::Sglang - } - } else { - RuntimeType::Sglang - }; - - // Build circuit breaker config - 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), - } - }; - - // Build health config - 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, - } - }; - - // Normalize URL: add protocol prefix only if missing - let normalized_url = if config.url.starts_with("http://") - || config.url.starts_with("https://") - || config.url.starts_with("grpc://") - { - config.url.clone() - } else { - match connection_mode.as_ref() { - ConnectionMode::Http => format!("http://{}", config.url), - ConnectionMode::Grpc { .. } => format!("grpc://{}", config.url), - } - }; - - if normalized_url != config.url { - debug!( - "Normalized worker URL: {} -> {} ({:?})", - config.url, - normalized_url, - connection_mode.as_ref() - ); - } - - // Handle DP-aware vs non-DP-aware workers - if config.dp_aware { - let dp_info: Arc = context.get_or_err("dp_info")?; - - debug!( - "Creating {} DP-aware workers for {} (dp_size: {})", - dp_info.dp_size, config.url, dp_info.dp_size - ); - - let mut workers = Vec::new(); - for rank in 0..dp_info.dp_size { - let mut builder = - DPAwareWorkerBuilder::new(normalized_url.clone(), rank, dp_info.dp_size) - .model(model_card.clone()) - .worker_type(worker_type.clone()) - .connection_mode(connection_mode.as_ref().clone()) - .runtime_type(runtime_type.clone()) - .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 !final_labels.is_empty() { - builder = builder.labels(final_labels.clone()); - } - - let worker = Arc::new(builder.build()) as Arc; - worker.set_healthy(false); - workers.push(worker); - - debug!( - "Created DP-aware worker {}@{}/{} ({:?})", - config.url, - rank, - dp_info.dp_size, - connection_mode.as_ref() - ); - } - - context.set("workers", workers); - context.set("labels", final_labels); - } else { - let mut builder = BasicWorkerBuilder::new(normalized_url.clone()) - .model(model_card) - .worker_type(worker_type) - .connection_mode(connection_mode.as_ref().clone()) - .runtime_type(runtime_type) - .circuit_breaker_config(circuit_breaker_config) - .health_config(health_config); - - 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()); - } - - let worker = Arc::new(builder.build()) as Arc; - worker.set_healthy(false); - - debug!( - "Created worker object for {} ({:?}) with {} labels", - config.url, - connection_mode.as_ref(), - final_labels.len() - ); - - context.set("worker", worker); - context.set("labels", final_labels); - } - - Ok(StepResult::Success) - } - - fn is_retryable(&self, _error: &WorkflowError) -> bool { - false - } -} - -/// Step 4: Register worker(s) in registry -pub struct RegisterWorkerStep; - -#[async_trait] -impl StepExecutor for RegisterWorkerStep { - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context.get_or_err("worker_config")?; - let app_context: Arc = context.get_or_err("app_context")?; - - if config.dp_aware { - let workers: Arc>> = context.get_or_err("workers")?; - - 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 DP-aware worker {} with ID {:?}", - config.url, worker_id - ); - } - - context.set("worker_ids", worker_ids); - } else { - let worker: Arc> = context.get_or_err("worker")?; - - 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) - } - - fn is_retryable(&self, _error: &WorkflowError) -> bool { - false - } -} - -/// Step 5: Update policy registry with worker information -pub struct UpdatePoliciesStep; - -#[async_trait] -impl StepExecutor for UpdatePoliciesStep { - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context.get_or_err("worker_config")?; - let labels: Arc> = context.get_or_err("labels")?; - let app_context: Arc = context.get_or_err("app_context")?; - - let policy_hint = labels.get("policy").map(|s| s.as_str()); - - if config.dp_aware { - let workers: Arc>> = context.get_or_err("workers")?; - - let model_id = workers[0].model_id().to_string(); - - for _ in 0..workers.len() { - app_context - .policy_registry - .on_worker_added(&model_id, policy_hint); - } - - // Initialize cache-aware policy if needed - let all_workers = app_context.worker_registry.get_by_model_fast(&model_id); - if let Some(policy) = app_context.policy_registry.get_policy(&model_id) { - if policy.name() == "cache_aware" { - app_context - .policy_registry - .init_cache_aware_policy(&model_id, &all_workers); - } - } - - debug!( - "Updated policies for {} DP-aware workers {} (model: {})", - workers.len(), - config.url, - model_id - ); - } else { - let worker: Arc> = context.get_or_err("worker")?; - - let model_id = worker.model_id().to_string(); - - app_context - .policy_registry - .on_worker_added(&model_id, policy_hint); - - // Initialize cache-aware policy if needed - let all_workers = app_context.worker_registry.get_by_model_fast(&model_id); - if let Some(policy) = app_context.policy_registry.get_policy(&model_id) { - if policy.name() == "cache_aware" { - app_context - .policy_registry - .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" { - app_context - .policy_registry - .init_pd_bucket_policies(&prefill_workers); - } - - debug!( - "Updated policies for worker {} (model: {})", - config.url, model_id - ); - } - - Ok(StepResult::Success) - } - - fn is_retryable(&self, _error: &WorkflowError) -> bool { - false - } -} - -/// Step 6: Activate worker(s) by marking them as healthy -pub struct ActivateWorkerStep; - -#[async_trait] -impl StepExecutor for ActivateWorkerStep { - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context.get_or_err("worker_config")?; - - if config.dp_aware { - let workers: Arc>> = context.get_or_err("workers")?; - - for worker in workers.iter() { - worker.set_healthy(true); - } - - info!( - "Activated {} DP-aware workers from {} (marked as healthy)", - workers.len(), - config.url - ); - } else { - let worker: Arc> = context.get_or_err("worker")?; - - worker.set_healthy(true); - - info!("Activated worker {} (marked as healthy)", config.url); - } - - Ok(StepResult::Success) - } - - fn is_retryable(&self, _error: &WorkflowError) -> bool { - false - } -} - -/// Create local worker registration workflow definition -pub fn create_worker_registration_workflow( - router_config: &crate::config::RouterConfig, -) -> WorkflowDefinition { - let detect_timeout = Duration::from_secs(router_config.worker_startup_timeout_secs); - - // 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 { - (5 + ((effective_timeout - 10.0) / 5.0).ceil() as u32).max(3) - } else { - 3 - }; - - WorkflowDefinition::new("worker_registration", "Worker Registration") - .add_step( - StepDefinition::new( - "detect_connection_mode", - "Detect Connection Mode", - Arc::new(DetectConnectionModeStep), - ) - .with_retry(RetryPolicy { - max_attempts, - backoff: BackoffStrategy::Linear { - increment: Duration::from_secs(1), - max: Duration::from_secs(5), - }, - }) - .with_timeout(detect_timeout) - .with_failure_action(FailureAction::FailWorkflow), - ) - .add_step( - StepDefinition::new( - "discover_metadata", - "Discover Metadata", - Arc::new(DiscoverMetadataStep), - ) - .with_retry(RetryPolicy { - max_attempts: 3, - backoff: BackoffStrategy::Fixed(Duration::from_secs(1)), - }) - .with_timeout(Duration::from_secs(10)) - .with_failure_action(FailureAction::ContinueNextStep) - .depends_on(&["detect_connection_mode"]), - ) - .add_step( - StepDefinition::new( - "discover_dp_info", - "Discover DP Info", - Arc::new(DiscoverDPInfoStep), - ) - .with_retry(RetryPolicy { - max_attempts: 3, - backoff: BackoffStrategy::Fixed(Duration::from_secs(1)), - }) - .with_timeout(Duration::from_secs(10)) - .with_failure_action(FailureAction::FailWorkflow) - .depends_on(&["discover_metadata"]), - ) - .add_step( - StepDefinition::new("create_worker", "Create Worker", Arc::new(CreateWorkerStep)) - .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::FailWorkflow) - .depends_on(&["discover_dp_info"]), - ) - .add_step( - StepDefinition::new( - "register_worker", - "Register Worker", - Arc::new(RegisterWorkerStep), - ) - .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::FailWorkflow) - .depends_on(&["create_worker"]), - ) - .add_step( - StepDefinition::new( - "update_policies", - "Update Policies", - Arc::new(UpdatePoliciesStep), - ) - .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::ContinueNextStep) - .depends_on(&["register_worker"]), - ) - .add_step( - StepDefinition::new( - "activate_worker", - "Activate Worker", - Arc::new(ActivateWorkerStep), - ) - .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::FailWorkflow) - .depends_on(&["update_policies"]), - ) -}