chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[model-gateway] add workflow for external model providers (#14323)

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