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

[model-gateway] Refactor worker steps and add update workflow (#15085)

Browse Source
This commit is contained in:
Simo Lin
2025-12-13 15:00:16 -08:00
committed by GitHub
parent 96705514bd
commit fd37cc5d38
14 changed files with 743 additions and 245 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -22,7 +22,7 @@ use crate::{
},
mcp::McpConfig,
observability::metrics::RouterMetrics,
protocols::worker_spec::{JobStatus, WorkerConfigRequest},
protocols::worker_spec::{JobStatus, WorkerConfigRequest, WorkerUpdateRequest},
workflow::{WorkflowContext, WorkflowEngine, WorkflowId, WorkflowInstanceId, WorkflowStatus},
};
@@ -32,6 +32,10 @@ pub enum Job {
AddWorker {
config: Box<WorkerConfigRequest>,
},
UpdateWorker {
url: String,
update: Box<WorkerUpdateRequest>,
},
RemoveWorker {
url: String,
},
@@ -57,6 +61,7 @@ impl Job {
pub fn job_type(&self) -> &str {
match self {
Job::AddWorker { .. } => "AddWorker",
Job::UpdateWorker { .. } => "UpdateWorker",
Job::RemoveWorker { .. } => "RemoveWorker",
Job::InitializeWorkersFromConfig { .. } => "InitializeWorkersFromConfig",
Job::InitializeMcpServers { .. } => "InitializeMcpServers",
@@ -70,6 +75,7 @@ impl Job {
pub fn worker_url(&self) -> &str {
match self {
Job::AddWorker { config } => &config.url,
Job::UpdateWorker { url, .. } => url,
Job::RemoveWorker { url } => url,
Job::InitializeWorkersFromConfig { .. } => "startup",
Job::InitializeMcpServers { .. } => "startup",
@@ -342,6 +348,24 @@ impl JobQueue {
)
.await
}
Job::UpdateWorker { url, update } => {
let engine = context
.workflow_engine
.get()
.ok_or_else(|| "Workflow engine not initialized".to_string())?;
let instance_id =
Self::start_worker_update_workflow(engine, url, update, context).await?;
debug!(
"Started worker update workflow for {} (instance: {})",
url, instance_id
);
let timeout_duration = Duration::from_secs(30);
Self::wait_for_workflow_completion(engine, instance_id, url, timeout_duration).await
}
Job::RemoveWorker { url } => {
let engine = context
.workflow_engine
@@ -684,6 +708,26 @@ impl JobQueue {
.map_err(|e| format!("Failed to start worker removal workflow: {:?}", e))
}
/// Start worker update workflow
async fn start_worker_update_workflow(
engine: &Arc<WorkflowEngine>,
url: &str,
update: &WorkerUpdateRequest,
context: &Arc<AppContext>,
) -> Result<WorkflowInstanceId, String> {
let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
// Pass URL and dp_aware separately, workflow step handles the rest
workflow_context.set("worker_url", url.to_string());
workflow_context.set("dp_aware", context.router_config.dp_aware);
workflow_context.set("update_request", update.clone());
workflow_context.set_arc("app_context", Arc::clone(context));
engine
.start_workflow(WorkflowId::new("worker_update"), workflow_context)
.await
.map_err(|e| format!("Failed to start worker update workflow: {:?}", e))
}
/// Start MCP server registration workflow
async fn start_mcp_registration_workflow(
engine: &Arc<WorkflowEngine>,

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

@@ -1,7 +1,7 @@
//! Workflow step implementations
//!
//! This module contains concrete step implementations for various workflows:
//! - Worker management (registration, removal, future: updates)
//! - Worker management (registration, removal, updates)
//! - MCP server registration
//! - WASM module registration and removal
//! - Future: Tokenizer fetching, LoRA updates, etc.
@@ -22,6 +22,7 @@ pub use worker::{
create_external_worker_workflow,
create_local_worker_workflow,
create_worker_removal_workflow,
create_worker_update_workflow,
// Utility functions
group_models_into_cards,
infer_model_type_from_id,
@@ -36,6 +37,8 @@ pub use worker::{
DiscoverMetadataStep,
DiscoverModelsStep,
DpInfo,
// Update steps
FindWorkerToUpdateStep,
// Removal steps
FindWorkersToRemoveStep,
ModelInfo,
@@ -43,8 +46,10 @@ pub use worker::{
RegisterWorkersStep,
RemoveFromPolicyRegistryStep,
RemoveFromWorkerRegistryStep,
UpdatePoliciesForWorkerStep,
UpdatePoliciesStep,
UpdateRemainingPoliciesStep,
UpdateWorkerPropertiesStep,
WorkerList,
WorkerRemovalRequest,
};

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

@@ -0,0 +1,64 @@
//! Step to find a worker to update based on URL.
use std::sync::Arc;
use async_trait::async_trait;
use tracing::debug;
use super::find_workers_by_url;
use crate::{
app_context::AppContext,
workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
/// Step to find workers to update based on URL.
///
/// For DP-aware workers, finds all workers with matching URL prefix.
/// For regular workers, finds the single worker with exact URL match.
///
/// Expects the following context values:
/// - "worker_url": String - the URL of the worker to update
/// - "dp_aware": bool - whether to find all DP-aware workers with matching prefix
/// - "app_context": Arc<AppContext>
///
/// Sets the following context values:
/// - "workers_to_update": Vec<Arc<dyn Worker>>
pub struct FindWorkerToUpdateStep;
#[async_trait]
impl StepExecutor for FindWorkerToUpdateStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let worker_url: Arc<String> = context.get_or_err("worker_url")?;
let dp_aware: Arc<bool> = context.get_or_err("dp_aware")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers_to_update =
find_workers_by_url(&app_context.worker_registry, &worker_url, *dp_aware);
if workers_to_update.is_empty() {
let error_msg = if *dp_aware {
format!("No workers found with prefix {}@", *worker_url)
} else {
format!("Worker {} not found", *worker_url)
};
return Err(WorkflowError::StepFailed {
step_id: StepId::new("find_worker_to_update"),
message: error_msg,
});
}
debug!(
"Found {} worker(s) to update for {}",
workers_to_update.len(),
*worker_url
);
context.set("workers_to_update", workers_to_update);
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false
}
}

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

@@ -0,0 +1,75 @@
//! Step to find workers to remove based on URL.
use std::{collections::HashSet, sync::Arc};
use async_trait::async_trait;
use tracing::debug;
use super::find_workers_by_url;
use crate::{
app_context::AppContext,
workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
/// Request structure for worker removal.
#[derive(Debug, Clone)]
pub struct WorkerRemovalRequest {
pub url: String,
pub dp_aware: bool,
}
/// Step to find workers to remove based on URL.
///
/// For DP-aware workers, finds all workers with matching URL prefix.
/// For regular workers, finds the single worker with exact URL match.
pub struct FindWorkersToRemoveStep;
#[async_trait]
impl StepExecutor for FindWorkersToRemoveStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let request: Arc<WorkerRemovalRequest> = context.get_or_err("removal_request")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers_to_remove =
find_workers_by_url(&app_context.worker_registry, &request.url, request.dp_aware);
if workers_to_remove.is_empty() {
let error_msg = if request.dp_aware {
format!("No workers found with prefix {}@", request.url)
} else {
format!("Worker {} not found", request.url)
};
return Err(WorkflowError::StepFailed {
step_id: StepId::new("find_workers_to_remove"),
message: error_msg,
});
}
debug!(
"Found {} worker(s) to remove for {}",
workers_to_remove.len(),
request.url
);
// Store workers and their model IDs for subsequent steps
let worker_urls: Vec<String> = workers_to_remove
.iter()
.map(|w| w.url().to_string())
.collect();
let affected_models: HashSet<String> = workers_to_remove
.iter()
.map(|w| w.model_id().to_string())
.collect();
context.set("workers_to_remove", workers_to_remove);
context.set("worker_urls", worker_urls);
context.set("affected_models", affected_models);
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false
}
}

97
sgl-model-gateway/src/core/steps/worker/local/mod.rs
View File

@@ -2,7 +2,13 @@ mod create_worker;
mod detect_connection;
mod discover_dp;
mod discover_metadata;
mod removal;
mod find_worker_to_update;
mod find_workers_to_remove;
mod remove_from_policy_registry;
mod remove_from_worker_registry;
mod update_policies_for_worker;
mod update_remaining_policies;
mod update_worker_properties;
use std::{sync::Arc, time::Duration};
@@ -10,17 +16,46 @@ 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,
};
pub use find_worker_to_update::FindWorkerToUpdateStep;
pub use find_workers_to_remove::{FindWorkersToRemoveStep, WorkerRemovalRequest};
pub use remove_from_policy_registry::RemoveFromPolicyRegistryStep;
pub use remove_from_worker_registry::RemoveFromWorkerRegistryStep;
pub use update_policies_for_worker::UpdatePoliciesForWorkerStep;
pub use update_remaining_policies::UpdateRemainingPoliciesStep;
pub use update_worker_properties::UpdateWorkerPropertiesStep;
use super::shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep};
use crate::{
config::RouterConfig,
core::{Worker, WorkerRegistry},
workflow::{BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition},
};
/// Find workers by URL, supporting both DP-aware (prefix match) and regular (exact match) modes.
///
/// For DP-aware workers, finds all workers with URL prefix `{url}@`.
/// For regular workers, finds the single worker with exact URL match.
pub(crate) fn find_workers_by_url(
registry: &WorkerRegistry,
url: &str,
dp_aware: bool,
) -> Vec<Arc<dyn Worker>> {
if dp_aware {
let worker_url_prefix = format!("{}@", url);
registry
.get_all()
.iter()
.filter(|worker| worker.url().starts_with(&worker_url_prefix))
.cloned()
.collect()
} else {
match registry.get_by_url(url) {
Some(worker) => vec![worker],
None => Vec::new(),
}
}
}
pub fn create_local_worker_workflow(router_config: &RouterConfig) -> WorkflowDefinition {
let detect_timeout = Duration::from_secs(router_config.worker_startup_timeout_secs);
@@ -193,3 +228,55 @@ pub fn create_worker_removal_workflow() -> WorkflowDefinition {
.depends_on(&["remove_from_worker_registry"]),
)
}
/// Create a worker update workflow definition.
///
/// DAG structure:
/// ```text
/// find_worker_to_update
/// │
/// update_worker_properties
/// │
/// update_policies_for_worker
/// ```
pub fn create_worker_update_workflow() -> WorkflowDefinition {
WorkflowDefinition::new("worker_update", "Update worker properties")
.add_step(
StepDefinition::new(
"find_worker_to_update",
"Find worker to update",
Arc::new(FindWorkerToUpdateStep),
)
.with_timeout(Duration::from_secs(10))
.with_retry(RetryPolicy {
max_attempts: 1,
backoff: BackoffStrategy::Fixed(Duration::from_secs(0)),
}),
)
.add_step(
StepDefinition::new(
"update_worker_properties",
"Update worker properties",
Arc::new(UpdateWorkerPropertiesStep),
)
.with_timeout(Duration::from_secs(10))
.with_retry(RetryPolicy {
max_attempts: 1,
backoff: BackoffStrategy::Fixed(Duration::from_secs(0)),
})
.depends_on(&["find_worker_to_update"]),
)
.add_step(
StepDefinition::new(
"update_policies_for_worker",
"Update policies for updated worker",
Arc::new(UpdatePoliciesForWorkerStep),
)
.with_timeout(Duration::from_secs(10))
.with_retry(RetryPolicy {
max_attempts: 1,
backoff: BackoffStrategy::Fixed(Duration::from_secs(0)),
})
.depends_on(&["update_worker_properties"]),
)
}

230
sgl-model-gateway/src/core/steps/worker/local/removal.rs
View File

@@ -1,230 +0,0 @@
//! 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::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
/// 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.
pub struct FindWorkersToRemoveStep;
#[async_trait]
impl StepExecutor for FindWorkersToRemoveStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let request: Arc<WorkerRemovalRequest> = context.get_or_err("removal_request")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
debug!(
"Finding workers to remove for {} (dp_aware: {})",
request.url, request.dp_aware
);
let workers_to_remove: Vec<Arc<dyn Worker>> = if request.dp_aware {
// DP-aware: Find all workers with matching prefix
let worker_url_prefix = format!("{}@", request.url);
let all_workers = app_context.worker_registry.get_all();
all_workers
.iter()
.filter(|worker| worker.url().starts_with(&worker_url_prefix))
.cloned()
.collect()
} else {
// Non-DP-aware: Find single worker by exact URL
match app_context.worker_registry.get_by_url(&request.url) {
Some(worker) => vec![worker],
None => Vec::new(),
}
};
if workers_to_remove.is_empty() {
let error_msg = if request.dp_aware {
format!("No workers found with prefix {}@", request.url)
} else {
format!("Worker {} not found", request.url)
};
return Err(WorkflowError::StepFailed {
step_id: StepId::new("find_workers_to_remove"),
message: error_msg,
});
}
debug!(
"Found {} worker(s) to remove for {}",
workers_to_remove.len(),
request.url
);
// Store workers and their model IDs for subsequent steps
let worker_urls: Vec<String> = workers_to_remove
.iter()
.map(|w| w.url().to_string())
.collect();
let affected_models: HashSet<String> = workers_to_remove
.iter()
.map(|w| w.model_id().to_string())
.collect();
context.set("workers_to_remove", workers_to_remove);
context.set("worker_urls", worker_urls);
context.set("affected_models", affected_models);
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false
}
}
/// Step 2: Remove workers from policy registry.
pub struct RemoveFromPolicyRegistryStep;
#[async_trait]
impl StepExecutor for RemoveFromPolicyRegistryStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers_to_remove: Arc<Vec<Arc<dyn Worker>>> =
context.get_or_err("workers_to_remove")?;
debug!(
"Removing {} worker(s) from policy registry",
workers_to_remove.len()
);
for worker in workers_to_remove.iter() {
let model_id = worker.model_id().to_string();
let worker_url = worker.url();
// Remove from cache-aware policy
app_context
.policy_registry
.remove_worker_from_cache_aware(&model_id, worker_url);
// Notify policy registry
app_context.policy_registry.on_worker_removed(&model_id);
debug!(
"Removed worker {} from policy registry (model: {})",
worker_url, model_id
);
}
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false
}
}
/// Step 3: Remove workers from worker registry.
pub struct RemoveFromWorkerRegistryStep;
#[async_trait]
impl StepExecutor for RemoveFromWorkerRegistryStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
debug!(
"Removing {} worker(s) from worker registry",
worker_urls.len()
);
let mut removed_count = 0;
for worker_url in worker_urls.iter() {
if app_context
.worker_registry
.remove_by_url(worker_url)
.is_some()
{
removed_count += 1;
debug!("Removed worker {} from registry", worker_url);
}
}
if removed_count != worker_urls.len() {
return Err(WorkflowError::StepFailed {
step_id: StepId::new("remove_from_worker_registry"),
message: format!(
"Expected to remove {} workers but only removed {}",
worker_urls.len(),
removed_count
),
});
}
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false
}
}
/// Step 4: Update cache-aware policies for remaining workers.
pub struct UpdateRemainingPoliciesStep;
#[async_trait]
impl StepExecutor for UpdateRemainingPoliciesStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let affected_models: Arc<HashSet<String>> = context.get_or_err("affected_models")?;
let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
debug!(
"Updating cache-aware policies for {} affected model(s)",
affected_models.len()
);
for model_id in affected_models.iter() {
let remaining_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" && !remaining_workers.is_empty() {
app_context
.policy_registry
.init_cache_aware_policy(model_id, &remaining_workers);
debug!(
"Updated cache-aware policy for model {} ({} remaining workers)",
model_id,
remaining_workers.len()
);
}
}
}
// Log final result at info level
if worker_urls.len() == 1 {
info!("Removed worker {}", worker_urls[0]);
} else {
info!(
"Removed {} DP-aware workers: {:?}",
worker_urls.len(),
worker_urls
);
}
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false
}
}

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

@@ -0,0 +1,56 @@
//! Step to remove workers from policy registry.
use std::sync::Arc;
use async_trait::async_trait;
use tracing::debug;
use crate::{
app_context::AppContext,
core::Worker,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
/// Step to remove workers from the policy registry.
///
/// Removes each worker from cache-aware policies and notifies
/// the policy registry of worker removal.
pub struct RemoveFromPolicyRegistryStep;
#[async_trait]
impl StepExecutor for RemoveFromPolicyRegistryStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers_to_remove: Arc<Vec<Arc<dyn Worker>>> =
context.get_or_err("workers_to_remove")?;
debug!(
"Removing {} worker(s) from policy registry",
workers_to_remove.len()
);
for worker in workers_to_remove.iter() {
let model_id = worker.model_id().to_string();
let worker_url = worker.url();
// Remove from cache-aware policy
app_context
.policy_registry
.remove_worker_from_cache_aware(&model_id, worker_url);
// Notify policy registry
app_context.policy_registry.on_worker_removed(&model_id);
}
debug!(
"Removed {} worker(s) from policy registry",
workers_to_remove.len()
);
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false
}
}

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

@@ -0,0 +1,57 @@
//! Step to remove workers from worker registry.
use std::sync::Arc;
use async_trait::async_trait;
use tracing::{debug, warn};
use crate::{
app_context::AppContext,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
/// Step to remove workers from the worker registry.
///
/// Removes each worker by URL from the central worker registry.
pub struct RemoveFromWorkerRegistryStep;
#[async_trait]
impl StepExecutor for RemoveFromWorkerRegistryStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
debug!(
"Removing {} worker(s) from worker registry",
worker_urls.len()
);
let mut removed_count = 0;
for worker_url in worker_urls.iter() {
if app_context
.worker_registry
.remove_by_url(worker_url)
.is_some()
{
removed_count += 1;
}
}
// Log if some workers were already removed (e.g., by another process)
if removed_count != worker_urls.len() {
warn!(
"Removed {} of {} workers (some may have been removed by another process)",
removed_count,
worker_urls.len()
);
} else {
debug!("Removed {} worker(s) from registry", removed_count);
}
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false
}
}

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

@@ -0,0 +1,65 @@
//! Step to update policies for updated workers.
use std::{collections::HashSet, sync::Arc};
use async_trait::async_trait;
use tracing::debug;
use crate::{
app_context::AppContext,
core::Worker,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
/// Step to update policies for updated workers.
///
/// After workers are updated, this step re-initializes cache-aware policies
/// for the affected models to reflect any priority or label changes.
pub struct UpdatePoliciesForWorkerStep;
#[async_trait]
impl StepExecutor for UpdatePoliciesForWorkerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let updated_workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("updated_workers")?;
// Collect affected models
let affected_models: HashSet<String> = updated_workers
.iter()
.map(|w| w.model_id().to_string())
.collect();
debug!(
"Updating policies for {} affected model(s) after worker update",
affected_models.len()
);
for model_id in &affected_models {
let 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" && !workers.is_empty() {
// Re-initialize cache-aware policy with updated workers
app_context
.policy_registry
.init_cache_aware_policy(model_id, &workers);
debug!(
"Updated cache-aware policy for model {} ({} workers)",
model_id,
workers.len()
);
}
}
// Notify policy registry of the update
app_context.policy_registry.on_worker_added(model_id, None);
}
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false
}
}

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

@@ -0,0 +1,66 @@
//! Step to update cache-aware policies for remaining workers after removal.
use std::{collections::HashSet, sync::Arc};
use async_trait::async_trait;
use tracing::{debug, info};
use crate::{
app_context::AppContext,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
/// Step to update cache-aware policies for remaining workers.
///
/// After workers are removed, this step re-initializes cache-aware policies
/// for the affected models using the remaining workers.
pub struct UpdateRemainingPoliciesStep;
#[async_trait]
impl StepExecutor for UpdateRemainingPoliciesStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let affected_models: Arc<HashSet<String>> = context.get_or_err("affected_models")?;
let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
debug!(
"Updating cache-aware policies for {} affected model(s)",
affected_models.len()
);
for model_id in affected_models.iter() {
let remaining_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" && !remaining_workers.is_empty() {
app_context
.policy_registry
.init_cache_aware_policy(model_id, &remaining_workers);
debug!(
"Updated cache-aware policy for model {} ({} remaining workers)",
model_id,
remaining_workers.len()
);
}
}
}
// Log final result at info level
if worker_urls.len() == 1 {
info!("Removed worker {}", worker_urls[0]);
} else {
info!(
"Removed {} DP-aware workers: {:?}",
worker_urls.len(),
worker_urls
);
}
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false
}
}

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

@@ -0,0 +1,147 @@
//! Step to update worker properties.
use std::sync::Arc;
use async_trait::async_trait;
use tracing::{debug, info};
use crate::{
app_context::AppContext,
core::{BasicWorkerBuilder, HealthConfig, Worker},
protocols::worker_spec::WorkerUpdateRequest,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
/// Step to update worker properties.
///
/// This step creates new worker instances with updated properties and
/// re-registers them to replace the old workers in the registry.
///
/// Expects the following context values:
/// - "update_request": WorkerUpdateRequest (from protocols::worker_spec)
/// - "app_context": Arc<AppContext>
/// - "workers_to_update": Vec<Arc<dyn Worker>>
///
/// Sets the following context values:
/// - "updated_workers": Vec<Arc<dyn Worker>>
pub struct UpdateWorkerPropertiesStep;
#[async_trait]
impl StepExecutor for UpdateWorkerPropertiesStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let request: Arc<WorkerUpdateRequest> = context.get_or_err("update_request")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers_to_update: Arc<Vec<Arc<dyn Worker>>> =
context.get_or_err("workers_to_update")?;
debug!(
"Updating properties for {} worker(s)",
workers_to_update.len()
);
let mut updated_workers: Vec<Arc<dyn Worker>> = Vec::with_capacity(workers_to_update.len());
for worker in workers_to_update.iter() {
// Build updated labels - merge new labels into existing ones
let mut updated_labels = worker.metadata().labels.clone();
if let Some(ref new_labels) = request.labels {
for (key, value) in new_labels {
updated_labels.insert(key.clone(), value.clone());
}
}
// Update priority if specified (stored in labels)
if let Some(priority) = request.priority {
updated_labels.insert("priority".to_string(), priority.to_string());
}
// Update cost if specified (stored in labels)
if let Some(cost) = request.cost {
updated_labels.insert("cost".to_string(), cost.to_string());
}
// Build updated health config
let existing_health = &worker.metadata().health_config;
let updated_health_config = HealthConfig {
timeout_secs: request
.health_check_timeout_secs
.unwrap_or(existing_health.timeout_secs),
check_interval_secs: request
.health_check_interval_secs
.unwrap_or(existing_health.check_interval_secs),
endpoint: existing_health.endpoint.clone(),
failure_threshold: request
.health_failure_threshold
.unwrap_or(existing_health.failure_threshold),
success_threshold: request
.health_success_threshold
.unwrap_or(existing_health.success_threshold),
};
// Determine API key: use new one if provided, otherwise keep existing
let updated_api_key = request
.api_key
.clone()
.or_else(|| worker.metadata().api_key.clone());
// Create a new worker with updated properties
let new_worker: Arc<dyn Worker> = if worker.is_dp_aware() {
// For DP-aware workers, extract DP info and rebuild
let dp_rank = worker.dp_rank().unwrap_or(0);
let dp_size = worker.dp_size().unwrap_or(1);
let base_url = worker.base_url().to_string();
let mut builder =
crate::core::DPAwareWorkerBuilder::new(base_url, dp_rank, dp_size)
.worker_type(worker.worker_type().clone())
.connection_mode(worker.connection_mode().clone())
.runtime_type(worker.metadata().runtime_type.clone())
.labels(updated_labels)
.health_config(updated_health_config.clone())
.models(worker.metadata().models.clone());
if let Some(ref api_key) = updated_api_key {
builder = builder.api_key(api_key.clone());
}
Arc::new(builder.build())
} else {
// For basic workers, rebuild with updated properties
let mut builder = BasicWorkerBuilder::new(worker.url())
.worker_type(worker.worker_type().clone())
.connection_mode(worker.connection_mode().clone())
.runtime_type(worker.metadata().runtime_type.clone())
.labels(updated_labels)
.health_config(updated_health_config.clone())
.models(worker.metadata().models.clone());
if let Some(ref api_key) = updated_api_key {
builder = builder.api_key(api_key.clone());
}
Arc::new(builder.build())
};
// Re-register the worker (this replaces the old one)
app_context.worker_registry.register(new_worker.clone());
updated_workers.push(new_worker);
}
// Log result
if updated_workers.len() == 1 {
info!("Updated worker {}", updated_workers[0].url());
} else {
info!("Updated {} workers", updated_workers.len());
}
// Store updated workers for subsequent steps
context.set("updated_workers", updated_workers);
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false
}
}

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

@@ -9,9 +9,10 @@ pub use external::{
};
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,
create_local_worker_workflow, create_worker_removal_workflow, create_worker_update_workflow,
CreateLocalWorkerStep, DetectConnectionModeStep, DiscoverDPInfoStep, DiscoverMetadataStep,
DpInfo, FindWorkerToUpdateStep, FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep,
RemoveFromWorkerRegistryStep, UpdatePoliciesForWorkerStep, UpdateRemainingPoliciesStep,
UpdateWorkerPropertiesStep, WorkerRemovalRequest,
};
pub use shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep, WorkerList};

22
sgl-model-gateway/src/protocols/worker_spec.rs
View File

@@ -210,7 +210,7 @@ pub struct WorkerTypeStats {
}
/// Worker update request
#[derive(Debug, Clone, Deserialize)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkerUpdateRequest {
/// Update priority
#[serde(skip_serializing_if = "Option::is_none")]
@@ -223,6 +223,26 @@ pub struct WorkerUpdateRequest {
/// Update labels
#[serde(skip_serializing_if = "Option::is_none")]
pub labels: Option<HashMap<String, String>>,
/// Update API key (for key rotation)
#[serde(skip_serializing_if = "Option::is_none")]
pub api_key: Option<String>,
/// Update health check timeout in seconds
#[serde(skip_serializing_if = "Option::is_none")]
pub health_check_timeout_secs: Option<u64>,
/// Update health check interval in seconds
#[serde(skip_serializing_if = "Option::is_none")]
pub health_check_interval_secs: Option<u64>,
/// Update health success threshold
#[serde(skip_serializing_if = "Option::is_none")]
pub health_success_threshold: Option<u32>,
/// Update health failure threshold
#[serde(skip_serializing_if = "Option::is_none")]
pub health_failure_threshold: Option<u32>,
}
/// Generic API response

47
sgl-model-gateway/src/server.rs
View File

@@ -27,6 +27,7 @@ use crate::{
create_external_worker_registration_workflow, create_mcp_registration_workflow,
create_wasm_module_registration_workflow, create_wasm_module_removal_workflow,
create_worker_registration_workflow, create_worker_removal_workflow,
create_worker_update_workflow,
},
worker_to_info, Job, JobQueue, JobQueueConfig, WorkerManager, WorkerType,
},
@@ -45,7 +46,7 @@ use crate::{
rerank::{RerankRequest, V1RerankReqInput},
responses::{ResponsesGetParams, ResponsesRequest},
validated::ValidatedJson,
worker_spec::{WorkerConfigRequest, WorkerErrorResponse, WorkerInfo},
worker_spec::{WorkerConfigRequest, WorkerErrorResponse, WorkerInfo, WorkerUpdateRequest},
},
routers::{conversations, router_manager::RouterManager, RouterTrait},
service_discovery::{start_service_discovery, ServiceDiscoveryConfig},
@@ -588,6 +589,41 @@ async fn delete_worker(State(state): State<Arc<AppState>>, Path(url): Path<Strin
}
}
async fn update_worker(
State(state): State<Arc<AppState>>,
Path(url): Path<String>,
Json(update): Json<WorkerUpdateRequest>,
) -> Response {
let worker_id = url.clone();
let job = Job::UpdateWorker {
url,
update: Box::new(update),
};
let job_queue = state
.context
.worker_job_queue
.get()
.expect("JobQueue not initialized");
match job_queue.submit(job).await {
Ok(_) => {
let response = json!({
"status": "accepted",
"worker_id": worker_id,
"message": "Worker update queued for background processing"
});
(StatusCode::ACCEPTED, Json(response)).into_response()
}
Err(error) => {
let error_response = WorkerErrorResponse {
error,
code: "INTERNAL_SERVER_ERROR".to_string(),
};
(StatusCode::INTERNAL_SERVER_ERROR, Json(error_response)).into_response()
}
}
}
pub struct ServerConfig {
pub host: String,
pub port: u16,
@@ -679,8 +715,10 @@ pub fn build_app(
let worker_routes = Router::new()
.route("/workers", post(create_worker))
.route("/workers", get(list_workers_rest))
.route("/workers/{url}", get(get_worker))
.route("/workers/{url}", delete(delete_worker))
.route(
"/workers/{url}",
get(get_worker).put(update_worker).delete(delete_worker),
)
.route_layer(axum::middleware::from_fn_with_state(
auth_config.clone(),
middleware::auth_middleware,
@@ -779,6 +817,9 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
engine
.register_workflow(create_worker_removal_workflow())
.expect("worker_removal workflow should be valid");
engine
.register_workflow(create_worker_update_workflow())
.expect("worker_update workflow should be valid");
engine
.register_workflow(create_mcp_registration_workflow())
.expect("mcp_registration workflow should be valid");