diff --git a/sgl-model-gateway/src/core/job_queue.rs b/sgl-model-gateway/src/core/job_queue.rs index 06261a6cd..9aebc0d86 100644 --- a/sgl-model-gateway/src/core/job_queue.rs +++ b/sgl-model-gateway/src/core/job_queue.rs @@ -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, }, + UpdateWorker { + url: String, + update: Box, + }, 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, + url: &str, + update: &WorkerUpdateRequest, + context: &Arc, + ) -> Result { + 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, diff --git a/sgl-model-gateway/src/core/steps/mod.rs b/sgl-model-gateway/src/core/steps/mod.rs index e19165c48..9642a0586 100644 --- a/sgl-model-gateway/src/core/steps/mod.rs +++ b/sgl-model-gateway/src/core/steps/mod.rs @@ -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, }; diff --git a/sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs b/sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs new file mode 100644 index 000000000..2bf9d154d --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs @@ -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 +/// +/// Sets the following context values: +/// - "workers_to_update": Vec> +pub struct FindWorkerToUpdateStep; + +#[async_trait] +impl StepExecutor for FindWorkerToUpdateStep { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + let worker_url: Arc = context.get_or_err("worker_url")?; + let dp_aware: Arc = context.get_or_err("dp_aware")?; + let app_context: Arc = 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 + } +} diff --git a/sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs b/sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs new file mode 100644 index 000000000..72441b60a --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs @@ -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 { + let request: Arc = context.get_or_err("removal_request")?; + let app_context: Arc = 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 = workers_to_remove + .iter() + .map(|w| w.url().to_string()) + .collect(); + + let affected_models: HashSet = 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 + } +} diff --git a/sgl-model-gateway/src/core/steps/worker/local/mod.rs b/sgl-model-gateway/src/core/steps/worker/local/mod.rs index 3a2890219..47a8a0e48 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/mod.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/mod.rs @@ -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> { + 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"]), + ) +} diff --git a/sgl-model-gateway/src/core/steps/worker/local/removal.rs b/sgl-model-gateway/src/core/steps/worker/local/removal.rs deleted file mode 100644 index d94e243ee..000000000 --- a/sgl-model-gateway/src/core/steps/worker/local/removal.rs +++ /dev/null @@ -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 { - let request: Arc = context.get_or_err("removal_request")?; - let app_context: Arc = context.get_or_err("app_context")?; - - debug!( - "Finding workers to remove for {} (dp_aware: {})", - request.url, request.dp_aware - ); - - let workers_to_remove: Vec> = 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 = workers_to_remove - .iter() - .map(|w| w.url().to_string()) - .collect(); - - let affected_models: HashSet = 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 { - let app_context: Arc = context.get_or_err("app_context")?; - let workers_to_remove: Arc>> = - 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 { - let app_context: Arc = context.get_or_err("app_context")?; - let worker_urls: Arc> = 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 { - let app_context: Arc = context.get_or_err("app_context")?; - let affected_models: Arc> = context.get_or_err("affected_models")?; - let worker_urls: Arc> = 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 - } -} diff --git a/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs b/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs new file mode 100644 index 000000000..2755abefd --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs @@ -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 { + let app_context: Arc = context.get_or_err("app_context")?; + let workers_to_remove: Arc>> = + 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 + } +} diff --git a/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs b/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs new file mode 100644 index 000000000..d9bd39012 --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs @@ -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 { + let app_context: Arc = context.get_or_err("app_context")?; + let worker_urls: Arc> = 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 + } +} diff --git a/sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs b/sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs new file mode 100644 index 000000000..2deeeae4c --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs @@ -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 { + let app_context: Arc = context.get_or_err("app_context")?; + let updated_workers: Arc>> = context.get_or_err("updated_workers")?; + + // Collect affected models + let affected_models: HashSet = 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 + } +} diff --git a/sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs b/sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs new file mode 100644 index 000000000..a02442c51 --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs @@ -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 { + let app_context: Arc = context.get_or_err("app_context")?; + let affected_models: Arc> = context.get_or_err("affected_models")?; + let worker_urls: Arc> = 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 + } +} diff --git a/sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs b/sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs new file mode 100644 index 000000000..8bd290d22 --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs @@ -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 +/// - "workers_to_update": Vec> +/// +/// Sets the following context values: +/// - "updated_workers": Vec> +pub struct UpdateWorkerPropertiesStep; + +#[async_trait] +impl StepExecutor for UpdateWorkerPropertiesStep { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + let request: Arc = context.get_or_err("update_request")?; + let app_context: Arc = context.get_or_err("app_context")?; + let workers_to_update: Arc>> = + context.get_or_err("workers_to_update")?; + + debug!( + "Updating properties for {} worker(s)", + workers_to_update.len() + ); + + let mut updated_workers: Vec> = 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 = 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 + } +} diff --git a/sgl-model-gateway/src/core/steps/worker/mod.rs b/sgl-model-gateway/src/core/steps/worker/mod.rs index dd8b1fc8a..7e3d1a669 100644 --- a/sgl-model-gateway/src/core/steps/worker/mod.rs +++ b/sgl-model-gateway/src/core/steps/worker/mod.rs @@ -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}; diff --git a/sgl-model-gateway/src/protocols/worker_spec.rs b/sgl-model-gateway/src/protocols/worker_spec.rs index 753b84db5..2d9a2a58e 100644 --- a/sgl-model-gateway/src/protocols/worker_spec.rs +++ b/sgl-model-gateway/src/protocols/worker_spec.rs @@ -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>, + + /// Update API key (for key rotation) + #[serde(skip_serializing_if = "Option::is_none")] + pub api_key: Option, + + /// Update health check timeout in seconds + #[serde(skip_serializing_if = "Option::is_none")] + pub health_check_timeout_secs: Option, + + /// Update health check interval in seconds + #[serde(skip_serializing_if = "Option::is_none")] + pub health_check_interval_secs: Option, + + /// Update health success threshold + #[serde(skip_serializing_if = "Option::is_none")] + pub health_success_threshold: Option, + + /// Update health failure threshold + #[serde(skip_serializing_if = "Option::is_none")] + pub health_failure_threshold: Option, } /// Generic API response diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index 5f992d756..de1874bb3 100644 --- a/sgl-model-gateway/src/server.rs +++ b/sgl-model-gateway/src/server.rs @@ -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>, Path(url): Path>, + Path(url): Path, + Json(update): Json, +) -> 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