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

[model-gateway] refactor: unify worker management into modular workflow structure (#15010)

Browse Source
This commit is contained in:
Simo Lin
2025-12-12 15:25:35 -08:00
committed by GitHub
parent b737a125c5
commit 9a56273ad2
18 changed files with 1796 additions and 1622 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
sgl-model-gateway/src/core/job_queue.rs
View File

@@ -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

591
sgl-model-gateway/src/core/steps/external_worker_registration.rs
View File

@@ -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<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_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::<Vec<ModelCard>>("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::<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_or_err("worker_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let model_cards: Arc<Vec<ModelCard>> = 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<dyn Worker>;
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<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_or_err("worker_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers: Arc<Vec<Arc<dyn Worker>>> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let labels: Arc<HashMap<String, String>> = context.get_or_err("labels")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers: Arc<Vec<Arc<dyn Worker>>> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let workers: Arc<Vec<Arc<dyn Worker>>> = 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"]),
)
}

67
sgl-model-gateway/src/core/steps/mod.rs
View File

@@ -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,
};

153
sgl-model-gateway/src/core/steps/worker/external/create_workers.rs vendored Normal file
View File

@@ -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<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let model_cards: Arc<Vec<ModelCard>> = 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<String, String> = 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<dyn Worker>;
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<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
}
}

273
sgl-model-gateway/src/core/steps/worker/external/discover_models.rs vendored Normal file
View File

@@ -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<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"));
/// 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>,
}
/// 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.
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<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
}
/// 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
);
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<StepResult> {
let config: Arc<WorkerConfigRequest> = 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::<Vec<ModelCard>>("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::<Vec<_>>()
);
context.set("model_cards", model_cards);
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &WorkflowError) -> bool {
true
}
}

104
sgl-model-gateway/src/core/steps/worker/external/mod.rs vendored Normal file
View File

@@ -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"]),
)
}

340
sgl-model-gateway/src/core/steps/worker/local/create_worker.rs Normal file
View File

@@ -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<Arc<dyn Worker>>` for downstream steps
pub struct CreateLocalWorkerStep;
#[async_trait]
impl StepExecutor for CreateLocalWorkerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?;
let discovered_labels: Arc<HashMap<String, String>> =
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<String, String>,
) -> 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::<Vec<String>>(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::<String>("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<String, String>,
) -> Result<Vec<Arc<dyn Worker>>, WorkflowError> {
let dp_info: Arc<DpInfo> = 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<dyn Worker>;
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<String, String>,
) -> Vec<Arc<dyn Worker>> {
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<dyn Worker>;
worker.set_healthy(false);
debug!(
"Created worker object for {} ({:?}) with {} labels",
normalized_url,
connection_mode,
final_labels.len()
);
vec![worker]
}

146
sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs Normal file
View File

@@ -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<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = 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
}
}

79
sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs Normal file
View File

@@ -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<DpInfo, String> {
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<StepResult> {
let config: Arc<WorkerConfigRequest> = 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
}
}

244
sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs Normal file
View File

@@ -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<Client> = 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<String>,
pub model_path: Option<String>,
pub served_model_name: Option<String>,
pub dp_size: Option<usize>,
pub version: Option<String>,
pub max_batch_size: Option<usize>,
pub max_total_tokens: Option<usize>,
pub max_prefill_tokens: Option<usize>,
pub max_running_requests: Option<usize>,
pub max_num_reqs: Option<usize>,
}
/// Model information returned from /model_info endpoint.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ModelInfo {
pub model_path: Option<String>,
pub tokenizer_path: Option<String>,
pub is_generation: Option<bool>,
pub model_type: Option<String>,
pub architectures: Option<Vec<String>>,
}
/// 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<ServerInfo, String> {
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::<Value>()
.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<ModelInfo, String> {
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::<ModelInfo>()
.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>, 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<HashMap<String, String>, 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let connection_mode: Arc<ConnectionMode> = 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
}
}

195
sgl-model-gateway/src/core/steps/worker/local/mod.rs Normal file
View File

@@ -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"]),
)
}

84
sgl-model-gateway/src/core/steps/worker_removal.rs → sgl-model-gateway/src/core/steps/worker/local/removal.rs
View File

@@ -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"]),
)
}

17
sgl-model-gateway/src/core/steps/worker/mod.rs Normal file
View File

@@ -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};

35
sgl-model-gateway/src/core/steps/worker/shared/activate.rs Normal file
View File

@@ -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<StepResult> {
let workers: Arc<Vec<Arc<dyn Worker>>> = 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
}
}

21
sgl-model-gateway/src/core/steps/worker/shared/mod.rs Normal file
View File

@@ -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<Arc<dyn Worker>>;

46
sgl-model-gateway/src/core/steps/worker/shared/register.rs Normal file
View File

@@ -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<Arc<dyn Worker>>`.
pub struct RegisterWorkersStep;
#[async_trait]
impl StepExecutor for RegisterWorkersStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers: Arc<Vec<Arc<dyn Worker>>> = 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
}
}

78
sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs Normal file
View File

@@ -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<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
let labels: Arc<HashMap<String, String>> = 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
}
}

943
sgl-model-gateway/src/core/steps/worker_registration.rs
View File

@@ -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<Client> = 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<String>,
model_path: Option<String>,
served_model_name: Option<String>,
dp_size: Option<usize>,
version: Option<String>,
max_batch_size: Option<usize>,
max_total_tokens: Option<usize>,
max_prefill_tokens: Option<usize>,
max_running_requests: Option<usize>,
max_num_reqs: Option<usize>,
}
/// Model information returned from /model_info endpoint
#[derive(Debug, Clone, Deserialize, Serialize)]
struct ModelInfo {
model_path: Option<String>,
tokenizer_path: Option<String>,
is_generation: Option<bool>,
/// HuggingFace model type string (e.g., "llama", "qwen2", "gpt_oss")
model_type: Option<String>,
/// Model architectures from HuggingFace config (e.g., ["LlamaForCausalLM"])
architectures: Option<Vec<String>>,
}
#[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<ServerInfo, String> {
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<ServerInfo, String> {
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::<Value>()
.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<ModelInfo, String> {
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::<ModelInfo>()
.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<DpInfo, String> {
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<HashMap<String, String>, 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>, 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let connection_mode: Arc<ConnectionMode> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?;
let discovered_labels: Arc<HashMap<String, String>> =
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::<Vec<String>>(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::<String>("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<DpInfo> = 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<dyn Worker>;
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<dyn Worker>;
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<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
if config.dp_aware {
let workers: Arc<Vec<Arc<dyn Worker>>> = 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<Arc<dyn Worker>> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let labels: Arc<HashMap<String, String>> = context.get_or_err("labels")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let policy_hint = labels.get("policy").map(|s| s.as_str());
if config.dp_aware {
let workers: Arc<Vec<Arc<dyn Worker>>> = 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<Arc<dyn Worker>> = 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<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
if config.dp_aware {
let workers: Arc<Vec<Arc<dyn Worker>>> = 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<Arc<dyn Worker>> = 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"]),
)
}